Hi,
I am trying to write a simple NN example in libtorch and I was able to successfully run this code on a CPU. However, after changing torch::Device to kCUDA, I was able to run the code on a GPU but it throws a seg fault error after finishing the run. Could someone please help me?
I am attaching my code below
#include<torch/torch.h>
#include<iostream>
#include<cmath>
#include<cstdio>
//The batch size for training
const int64_t N = 64;
//The input dimension
const int64_t D_in = 1000;
//The hidden dimension
const int64_t H = 100;
//The output dimension
const int64_t D_out = 10;
//Total number of steps
const int64_t tstep = 1000;
using namespace torch;
using namespace std;
struct TwoLayerNetImpl : nn::Module{
TwoLayerNetImpl(int D_in, int H, int D_out)
: linear1(nn::LinearOptions(D_in, H).bias(false)),
linear2(nn::LinearOptions(H, D_out).bias(false))
{
register_module("linear1", linear1);
register_module("linear2", linear2);
}
torch::Tensor forward(torch::Tensor x)
{
x = torch::clamp_min(linear1(x),0);
x = linear2(x);
return x;
}
nn::Linear linear1, linear2;
};
TORCH_MODULE(TwoLayerNet);
int main(int argc, char* argv[])
{
torch::manual_seed(1);
torch::Device device(torch::kCUDA); //change to kCPU to run on host
if (torch::cuda::is_available())
{
std::cout << "CUDA is available! Training on GPU" << std::endl;
}
torch::Tensor X = torch::randn({N, D_in}).to(device);
torch::Tensor Y = torch::randn({N, D_out}).to(device);
TwoLayerNet net(D_in,H,D_out);
net->to(device);
torch::optim::Adam optimizer(net->parameters(), torch::optim::AdamOptions(1e-4));
torch::nn::MSELoss criterion((torch::nn::MSELossOptions(torch::kSum)));
for(int64_t ts=0;ts<=tstep;++ts)
{
torch::Tensor Y_pred = net->forward(X);
torch::Tensor loss = criterion(Y_pred, Y);
if(ts % 100 == 0)
{
printf("\r[%4ld/%4ld] | D_loss: %e \n", ts,tstep,loss.item<float>());
}
optimizer.zero_grad();
loss.backward();
optimizer.step();
}
std::cout<<"Training complete!"<<std::endl;
return 0;
}
Thank you!