Floating point exception (core dumped) for UNet implementation
I am trying to do an implementation of KiuNet ( https://github.com/jeya-maria-jose/KiU-Net-pytorch ). But when I am executing the train command like so:
python train.py --train_dataset &quot;KiuNet/Train Folder/&quot; --val_dataset &quot;KiuNet/Validation Folder/&quot; --direc 'KiuNet/Results/' --batch_size 1 --epoch 200 --save_freq 10 --modelname &quot;kiunet&quot; --learning_rate 0.0001

I am getting the following error:
Traceback (most recent call last):
  File &quot;KiuNet/KiU-Net-pytorch/train.py&quot;, line 235, in &lt;module&gt;
    loss.backward()
  File &quot;/miniconda3/lib/python3.9/site-packages/torch/_tensor.py&quot;, line 487, in backward
    torch.autograd.backward(
  File &quot;/miniconda3/lib/python3.9/site-packages/torch/autograd/__init__.py&quot;, line 197, in backward
    Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
RuntimeError: cuDNN error: CUDNN_STATUS_INTERNAL_ERROR
../aten/src/ATen/native/cuda/NLLLoss2d.cu:104: nll_loss2d_forward_kernel: block: [0,0,0], thread: [847,0,0] Assertion `t &gt;= 0 &amp;&amp; t &lt; n_classes` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:104: nll_loss2d_forward_kernel: block: [0,0,0], thread: [958,0,0] Assertion `t &gt;= 0 &amp;&amp; t &lt; n_classes` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:104: nll_loss2d_forward_kernel: block: [0,0,0], thread: [703,0,0] Assertion `t &gt;= 0 &amp;&amp; t &lt; n_classes` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:104: nll_loss2d_forward_kernel: block: [0,0,0], thread: [830,0,0] Assertion `t &gt;= 0 &amp;&amp; t &lt; n_classes` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:104: nll_loss2d_forward_kernel: block: [0,0,0], thread: [831,0,0] Assertion `t &gt;= 0 &amp;&amp; t &lt; n_classes` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:104: nll_loss2d_forward_kernel: block: [0,0,0], thread: [575,0,0] Assertion `t &gt;= 0 &amp;&amp; t &lt; n_classes` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:104: nll_loss2d_forward_kernel: block: [0,0,0], thread: [974,0,0] Assertion `t &gt;= 0 &amp;&amp; t &lt; n_classes` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:104: nll_loss2d_forward_kernel: block: [0,0,0], thread: [77,0,0] Assertion `t &gt;= 0 &amp;&amp; t &lt; n_classes` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:104: nll_loss2d_forward_kernel: block: [0,0,0], thread: [78,0,0] Assertion `t &gt;= 0 &amp;&amp; t &lt; n_classes` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:104: nll_loss2d_forward_kernel: block: [0,0,0], thread: [719,0,0] Assertion `t &gt;= 0 &amp;&amp; t &lt; n_classes` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:104: nll_loss2d_forward_kernel: block: [0,0,0], thread: [720,0,0] Assertion `t &gt;= 0 &amp;&amp; t &lt; n_classes` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:104: nll_loss2d_forward_kernel: block: [0,0,0], thread: [592,0,0] Assertion `t &gt;= 0 &amp;&amp; t &lt; n_classes` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:104: nll_loss2d_forward_kernel: block: [0,0,0], thread: [593,0,0] Assertion `t &gt;= 0 &amp;&amp; t &lt; n_classes` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:104: nll_loss2d_forward_kernel: block: [0,0,0], thread: [209,0,0] Assertion `t &gt;= 0 &amp;&amp; t &lt; n_classes` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:104: nll_loss2d_forward_kernel: block: [0,0,0], thread: [465,0,0] Assertion `t &gt;= 0 &amp;&amp; t &lt; n_classes` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:104: nll_loss2d_forward_kernel: block: [0,0,0], thread: [337,0,0] Assertion `t &gt;= 0 &amp;&amp; t &lt; n_classes` failed.

When I am running the train command with CUDA_LAUNCH_BLOCKING=1 I get the following error:
../aten/src/ATen/native/cuda/NLLLoss2d.cu:104: nll_loss2d_forward_kernel: block: [0,0,0], thread: [840,0,0] Assertion `t &gt;= 0 &amp;&amp; t &lt; n_classes` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:104: nll_loss2d_forward_kernel: block: [0,0,0], thread: [580,0,0] Assertion `t &gt;= 0 &amp;&amp; t &lt; n_classes` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:104: nll_loss2d_forward_kernel: block: [0,0,0], thread: [453,0,0] Assertion `t &gt;= 0 &amp;&amp; t &lt; n_classes` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:104: nll_loss2d_forward_kernel: block: [0,0,0], thread: [326,0,0] Assertion `t &gt;= 0 &amp;&amp; t &lt; n_classes` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:104: nll_loss2d_forward_kernel: block: [0,0,0], thread: [71,0,0] Assertion `t &gt;= 0 &amp;&amp; t &lt; n_classes` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:104: nll_loss2d_forward_kernel: block: [0,0,0], thread: [712,0,0] Assertion `t &gt;= 0 &amp;&amp; t &lt; n_classes` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:104: nll_loss2d_forward_kernel: block: [0,0,0], thread: [198,0,0] Assertion `t &gt;= 0 &amp;&amp; t &lt; n_classes` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:104: nll_loss2d_forward_kernel: block: [0,0,0], thread: [199,0,0] Assertion `t &gt;= 0 &amp;&amp; t &lt; n_classes` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:104: nll_loss2d_forward_kernel: block: [0,0,0], thread: [968,0,0] Assertion `t &gt;= 0 &amp;&amp; t &lt; n_classes` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:104: nll_loss2d_forward_kernel: block: [0,0,0], thread: [959,0,0] Assertion `t &gt;= 0 &amp;&amp; t &lt; n_classes` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:104: nll_loss2d_forward_kernel: block: [0,0,0], thread: [830,0,0] Assertion `t &gt;= 0 &amp;&amp; t &lt; n_classes` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:104: nll_loss2d_forward_kernel: block: [0,0,0], thread: [574,0,0] Assertion `t &gt;= 0 &amp;&amp; t &lt; n_classes` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:104: nll_loss2d_forward_kernel: block: [0,0,0], thread: [702,0,0] Assertion `t &gt;= 0 &amp;&amp; t &lt; n_classes` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:104: nll_loss2d_forward_kernel: block: [0,0,0], thread: [191,0,0] Assertion `t &gt;= 0 &amp;&amp; t &lt; n_classes` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:104: nll_loss2d_forward_kernel: block: [0,0,0], thread: [318,0,0] Assertion `t &gt;= 0 &amp;&amp; t &lt; n_classes` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:104: nll_loss2d_forward_kernel: block: [0,0,0], thread: [319,0,0] Assertion `t &gt;= 0 &amp;&amp; t &lt; n_classes` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:104: nll_loss2d_forward_kernel: block: [0,0,0], thread: [446,0,0] Assertion `t &gt;= 0 &amp;&amp; t &lt; n_classes` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:104: nll_loss2d_forward_kernel: block: [0,0,0], thread: [63,0,0] Assertion `t &gt;= 0 &amp;&amp; t &lt; n_classes` failed.
Floating point exception (core dumped)

My torch and CUDA version are: '1.13.0+cu117'
My Python version: Python 3.9.12
Any help is much appreciated!

The repository author mentions the following.
&quot;This bug occurs when the ground truth masks have more classes than the number of classes in prediction. Please make sure you ground truth images have only 0 or 1 labels of pixels if you are training for binary segmentation. The datasets usually have the ground truth as 0 or 255 labels of pixels. So, please convert them to 0's and 1's.&quot;
