I could reproduce the issue. I think you have stumbled upon a pytorch bug in nll_loss2d.
I had played around with your example and I could notice that reduction='none' could be the reason. Error is triggered for other reduction (mean / sum) option.
Further, the reason could be that, the device assertion is not checked in case of reduction='none'.
  
      github.com
  
  
    pytorch/pytorch/blob/master/aten/src/THCUNN/generic/SpatialClassNLLCriterion.cu#L70-L99 1

if (reduction == at::Reduction::None) {  int64_t batch_size = THCTensor_(size)(state, input, 0);  int64_t H = THCTensor_(size)(state, input, 2);  int64_t W = THCTensor_(size)(state, input, 3);  int64_t count = batch_size * H * W;  THCTensor_(resize3d)(state, output, batch_size, H, W);  if (count == 0) {    // This guards from unnecessary operations and launching CUDA kernel with 0 blocks.    return;  }  if (weights) {    weights = THCTensor_(newContiguous)(state, weights);  }  SpatialClassNLLCriterion_updateOutput_no_reduce_kernel<scalar_t>    <<<GET_BLOCKS(count), CUDA_NUM_THREADS, 0, c10::cuda::getCurrentCUDAStream()>>>(      count,      toDeviceTensor<scalar_t, 4>(state, input),


  This file has been truncated. show original

  
  
    
    
  
  


THCudaCheck(cudaGetLastError()); is missing incase of reduction='none'.
I have raised a bug on behalf of you: nll_loss2d: t >= 0 && t < n_classes assertion is not checked when using GPU tensors and reduction='none' · Issue #49882 · pytorch/pytorch · GitHub 4