About RuntimeError: CUDA error: device-side assert triggered (Tensor move to GPU)

Question

Enviroment: Colab Pytorch version: 1.10.0+cu111

After running os.environ["CUDA_LAUNCH_BLOCKING"] = "1", the following Code worked on CPU but got ERROR on GPU

main_device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
Test_generator = DataLoader(dataset=Test_dataset,batch_size=60, shuffle=SHUFFLE)

for i, (X_val,y_val) in enumerate(Test_generator):
 if main_device == torch.device('cpu'):
  X_val, y_val = X_val.to(main_device).type(torch.FloatTensor), y_val.to(main_device).type(torch.FloatTensor)
 else:
  X_val, y_val = X_val.to(main_device).type(torch.cuda.FloatTensor), y_val.to(main_device).type(torch.cuda.FloatTensor)

Afterwards, I try to figure out why and check previous debuged answer. I have no problem about Label, so I test the following code

X_val.to(main_device)

---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
<ipython-input-35-9f9e092171e8> in <module>()
----> 1 X_val.to(main_device)

RuntimeError: CUDA error: device-side assert triggered

I think that it is main problem ... but I don't know why ? How can I fix it ? enter image description here

enter image description here