pytorch_loss_wrapper.h

Go to the documentation of this file.

#ifndef PYTORCH_LOSS_WRAPPER_H
#define PYTORCH_LOSS_WRAPPER_H
 
#include "cuberl/base/cubeai_config.h"
 
#ifdef USE_PYTORCH
 
#include "cuberl/base/cuberl_types.h"
#include "cuberl/utils/loss_type.h"
#include <torch/torch.h>
 
 
 
 
namespace cuberl {
namespace utils{
namespace pytorch {
    
    using namespace cubeai::utils;
 
class PyTorchLossWrapper
{
public:
 
    PyTorchLossWrapper(LossType type);
 
    torch_tensor_t calculate(torch_tensor_t input, torch_tensor_t target)const;
 
private:
 
    LossType type_;
 
};
 
}
}
}
 
#endif
#endif // PYTORCH_LOSS_WRAPPER_H