RuntimeError: shape '[2, 3, 9, 80, 80]' is invalid for input of size 230400

[Naeem Raza]

I am training for only 1 class and got this error.

Epoch gpu_mem box obj cls total targets img_size
0% 0/11 [00:08<?, ?it/s]
Traceback (most recent call last):
File "/content/drive/MyDrive/PhD-OBJ-1/PyTorch_YOLOv4-master/train.py", line 537, in
train(hyp, opt, device, tb_writer, wandb)
File "/content/drive/MyDrive/PhD-OBJ-1/PyTorch_YOLOv4-master/train.py", line 287, in train
pred = model(imgs) # forward
File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1501, in _call_impl
return forward_call(*args, **kwargs)
File "/content/drive/MyDrive/PhD-OBJ-1/PyTorch_YOLOv4-master/models/models.py", line 465, in forward
return self.forward_once(x)
File "/content/drive/MyDrive/PhD-OBJ-1/PyTorch_YOLOv4-master/models/models.py", line 518, in forward_once
yolo_out.append(module(x, out))
File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1501, in _call_impl
return forward_call(*args, **kwargs)
File "/content/drive/MyDrive/PhD-OBJ-1/PyTorch_YOLOv4-master/models/models.py", line 321, in forward
p = p.view(bs, self.na, self.no, self.ny, self.nx).permute(0, 1, 3, 4, 2).contiguous() # prediction
RuntimeError: shape '[2, 3, 9, 80, 80]' is invalid for input of size 230400