RuntimeError: weight tensor should be defined either for all 3 classes or no classes but got weight tensor of shape: [2]

[aishanii]

While running the main file I am getting this error.

----Loading dataset----
Dataset: MAHNOB-HCI
#Traning samples: 360
#Validation samples: 3
#Training distribution: [210 150]
wandb: Currently logged in as: vanshuagarwal11-03 (vanshuagarwal11-03-SRM Institute of Science and Technology). Use wandb login --relogin to force relogin
wandb: Tracking run with wandb version 0.17.4
wandb: Run data is saved locally in /content/wandb/run-20240708_111041-0am4ciyg
wandb: Run wandb offline to turn off syncing.
wandb: Syncing run proud-leaf-49
wandb: ⭐️ View project at https://wandb.ai/vanshuagarwal11-03-SRM%20Institute%20of%20Science%20and%20Technology/MHyEEG
wandb: 🚀 View run at https://wandb.ai/vanshuagarwal11-03-SRM%20Institute%20of%20Science%20and%20Technology/MHyEEG/runs/0am4ciyg
Number of parameters: 19663747

Running on GPU? True - gpu_num: 0
Train round: 0% 0/45 [00:00<?, ?batch/s]tensor([[ 0.0504, -0.0277, -0.0015],
[-0.1052, -0.0144, 0.0084],
[-0.1097, 0.0798, 0.0005],
[-0.0701, -0.0238, -0.0405],
[-0.0034, 0.0291, 0.0478],
[-0.0992, 0.0290, -0.0440],
[-0.1483, -0.0215, -0.0354],
[-0.0660, 0.0021, -0.0198]], device='cuda:0',
grad_fn=)
tensor([2, 1, 2, 2, 1, 1, 2, 1], device='cuda:0')
Traceback (most recent call last):
File "/content/drive/MyDrive/MHyEEG-main-share/MHyEEG-main/main.py", line 88, in
main(args, n_workers)
File "/content/drive/MyDrive/MHyEEG-main-share/MHyEEG-main/main.py", line 47, in main
trainer.train(train_loader, eval_loader)
File "/content/drive/MyDrive/MHyEEG-main-share/MHyEEG-main/training.py", line 92, in train
loss = self.criterion(outputs, labels)
File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1532, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1541, in _call_impl
return forward_call(*args, **kwargs)
File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/loss.py", line 1185, in forward
return F.cross_entropy(input, target, weight=self.weight,
File "/usr/local/lib/python3.10/dist-packages/torch/nn/functional.py", line 3086, in cross_entropy
return torch._C._nn.cross_entropy_loss(input, target, weight, _Reduction.get_enum(reduction), ignore_index, label_smoothing)
RuntimeError: weight tensor should be defined either for all 3 classes or no classes but got weight tensor of shape: [2]
wandb: 🚀 View run proud-leaf-49 at: https://wandb.ai/vanshuagarwal11-03-SRM%20Institute%20of%20Science%20and%20Technology/MHyEEG/runs/0am4ciyg
wandb: ⭐️ View project at: https://wandb.ai/vanshuagarwal11-03-SRM%20Institute%20of%20Science%20and%20Technology/MHyEEG
wandb: Synced 5 W&B file(s), 0 media file(s), 0 artifact file(s) and 0 other file(s)
wandb: Find logs at: ./wandb/run-20240708_111041-0am4ciyg/logs
wandb: WARNING The new W&B backend becomes opt-out in version 0.18.0; try it out with wandb.require("core")! See https://wandb.me/wandb-core for more information.