MobileNetV2SSD in Pytorch

I am working on a MobileNetV2 with SSD head for object detection in ipynb, VSCode. My code has a class for the dataset (which is in YOLO format with image size 640x640), a block for the model (class for Bottleneck, MobileNetV2, SSDHead, MobileNetV2SSD), function defined for iou and anchor matching, function defined for ssd_loss, function defined for custom_collate_fn, and training loop.

When I run the code, I get the following error. How do I fix it… I need simple explanation as I am beginner in the field.

File c:\Users\ca8mc\Desktop\netv2v8\netv2v8_env\lib\site-packages\torch\nn\functional.py:3104, in cross_entropy(input, target, weight, size_average, ignore_index, reduce, reduction, label_smoothing)
[3102] if size_average is not None or reduce is not None:
[3103] = _Reduction.legacy_get_string(size_average, reduce)
→ [3104] return torch._C._nn.cross_entropy_loss(input, target, weight, _Reduction.get_enum(reduction), ignore_index, label_smoothing)

ValueError: Expected input batch_size (4704) to match target batch_size (64).

Based on the error message you are using nn.CrossEntropyLoss which expects a model output and target with the same batch size (in dim0). Check why the batch sizes differ in your case and if you accidentally reshaped one tensor on a wrong way.