Code runs fine on CPU and GPU but gives seg fault at the end onGPU

neilmehta87: <p>Hi,</p>
<p>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?<br>
I am attaching my code below</p>
<pre><code class="lang-auto">#include&lt;torch/torch.h&gt;
#include&lt;iostream&gt;
#include&lt;cmath&gt;
#include&lt;cstdio&gt;

//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 &lt;&lt; "CUDA is available! Training on GPU" &lt;&lt; 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-&gt;to(device);
  
    torch::optim::Adam optimizer(net-&gt;parameters(), torch::optim::AdamOptions(1e-4));
    torch::nn::MSELoss criterion((torch::nn::MSELossOptions(torch::kSum)));
  
    for(int64_t ts=0;ts&lt;=tstep;++ts)
    {
      torch::Tensor Y_pred = net-&gt;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&lt;float&gt;());
      }
      optimizer.zero_grad();
      loss.backward();
      optimizer.step();
    }
    std::cout&lt;&lt;"Training complete!"&lt;&lt;std::endl;
   return 0;
}
</code></pre>
<p>Thank you!</p>

glaringlee: <p>Can you provide which libtorch version (including cuda version) you are using?<br>
What’s your cuda version?<br>
Can you paste your error msg?</p>

neilmehta87: <p>Hi,</p>
<p>I am using pytorch/v1.5.0-gpu. I downloaded the stable release of libtorch.<br>
I was using cuda 10.1.243 yesterday and it was giving me a segmentation fault error.<br>
However, after switching to cuda 10.2.89 and recompiling my code, I no longer get the error.<br>
I apologize for the false alarm.</p>
<p>Thank you</p>

glaringlee: <p><a class="mention" href="/u/neilmehta87">@neilmehta87</a></p>
<p>Glad to here that. One possible reason is that the libtorch you downloaded is for cuda 10.2.<br>
We have different libtorch binary for different version of cuda.<br>
You can find them here: <a href="https://pytorch.org/" rel="nofollow noopener">https://pytorch.org/</a></p>

neilmehta87: <p><a class="mention" href="/u/glaringlee">@glaringlee</a> Thank you</p>