transform = transforms.Compose([
    transforms.Resize((480, 640)),
    transforms.ToTensor(),
])

train_dataset = datasets.CocoDetection(
    root=args.train_dataset, annFile=args.train_annotation, transform=transform)

train_loader = DataLoader(train_dataset, batch_size=64)

for (img, anno) in train_loader:
    out = model(img)

python3.7/site-packages/torch/utils/data/_utils/collate.py", line 82, in default_collate
    raise RuntimeError('each element in list of batch should be of equal size')
RuntimeError: each element in list of batch should be of equal size