Thanks a lot for your suggestion!
I was making exactly the mistake you have pointed out here during normalization but when I ran the code it threw another error:
RuntimeError: multi-target not supported at C:/cb/pytorch_1000000000000/work/aten/src\THCUNN/generic/ClassNLLCriterion.cu:15
I checked other posts where you have mentioned one-hot is not applicable, mine are labelled images where images and labels are separate variables.
Full stack of the same is the following:
Traceback (most recent call last):
File “F:/Richa/Richa_Python_codes/Variational-Capsule-Routing/src/main_r.py”, line 136, in 
score = main(parser.parse_known_args()[0])
File “F:/Richa/Richa_Python_codes/Variational-Capsule-Routing/src/main_r.py”, line 107, in main
score = train(model, dataloaders, args)
File “F:\Richa\Richa_Python_codes\Variational-Capsule-Routing\src\train.py”, line 103, in train
test_loss, test_acc = evaluate(model, args, dataloaders[‘test’])
File “F:\Richa\Richa_Python_codes\Variational-Capsule-Routing\src\evaluate.py”, line 22, in evaluate
loss = F.cross_entropy(yhat, labels.cuda())
File “C:\Users\user.conda\envs\gpu-test4\lib\site-packages\torch\nn\functional.py”, line 2468, in cross_entropy
return nll_loss(log_softmax(input, 1), target, weight, None, ignore_index, None, reduction)
File “C:\Users\user.conda\envs\gpu-test4\lib\site-packages\torch\nn\functional.py”, line 2264, in nll_loss
ret = torch._C._nn.nll_loss(input, target, weight, _Reduction.get_enum(reduction), ignore_index)
RuntimeError: multi-target not supported at C:/cb/pytorch_1000000000000/work/aten/src\THCUNN/generic/ClassNLLCriterion.cu:15