Why I'm getting: CUDA error: device-side assert triggered - pytorch

Question

I have the following model:
==========================================================================================
Layer (type:depth-idx) Output Shape Param #
=====================================================================================
=====
Sequential -- --
├─Identity: 1-1 [64, 1, 6, 2] --
├─Conv2d: 1-2 [64, 64, 6, 2] 576
├─Flatten: 1-3 [64, 768] --
├─Linear: 1-4 [64, 10] 7,690
├─ReLU: 1-5 [64, 10] --
├─Linear: 1-6 [64, 1] 11
==========================================================================================
I'm trying to train the model and caluclate accuracy:
DEVICE = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
Loss = nn.CrossEntropyLoss()
inputs, labels = next(iter(oTrainDL))
inputs, labels = inputs.to(DEVICE), labels.to(DEVICE)
oModel = GetModel()
oModel.to(DEVICE)
oModel.train(True)
inputs = inputs.type(torch.cuda.FloatTensor)
mZ = oModel(inputs)
labels = labels.type(torch.cuda.LongTensor)
loss = Loss(mZ, labels)
with torch.no_grad():
print(f"{mZ.shape}, {labels.shape}")
print(f"{mZ.is_cuda}, {labels.is_cuda}")
res = (mZ == labels).float().mean().item()
But I'm getting an error:
RuntimeError Traceback (most recent call last)
Input In [13], in <cell line: 18>()
21 print(f"{mZ.shape}, {labels.shape}")
22 print(f"{mZ.is_cuda}, {labels.is_cuda}")
---> 23 res = (mZ == labels).float().mean().item()
RuntimeError: CUDA error: device-side assert triggered
CUDA kernel errors might be asynchronously reported at some other API call,so the stacktrace below might be incorrect.
For debugging consider passing CUDA_LAUNCH_BLOCKING=1.
What am I doing wrong ?
How can I fix it ?