This leads to the same error as if I change the model:
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/home/mschiappa/.local/lib/python3.6/site-packages/torch/nn/modules/module.py", line 541, in __call__
    result = self.forward(*input, **kwargs)
  File "/usr/lib/python3/dist-packages/torchvision/models/detection/generalized_rcnn.py", line 52, in forward
    proposals, proposal_losses = self.rpn(images, features, targets)
  File "/home/mschiappa/.local/lib/python3.6/site-packages/torch/nn/modules/module.py", line 541, in __call__
    result = self.forward(*input, **kwargs)
  File "/usr/lib/python3/dist-packages/torchvision/models/detection/rpn.py", line 411, in forward
    boxes, scores = self.filter_proposals(proposals, objectness, images.image_sizes, num_anchors_per_level)
  File "/usr/lib/python3/dist-packages/torchvision/models/detection/rpn.py", line 336, in filter_proposals
    keep = box_ops.batched_nms(boxes, scores, lvl, self.nms_thresh)
  File "/usr/lib/python3/dist-packages/torchvision/ops/boxes.py", line 72, in batched_nms
    keep = nms(boxes_for_nms, scores, iou_threshold)
  File "/usr/lib/python3/dist-packages/torchvision/ops/boxes.py", line 32, in nms
    _C = _lazy_import()
  File "/usr/lib/python3/dist-packages/torchvision/extension.py", line 12, in _lazy_import
    from torchvision import _C as C
ImportError: /usr/lib/python3/dist-packages/torchvision/_C.cpython-36m-x86_64-linux-gnu.so: undefined symbol: _ZN3c107Warning4warnENS_14SourceLocationENSt7__cxx1112basic_stringIcSt11char_traitsIcESaIcEEE

also that is what I want to do but when I use the dataset class and make my own custom getitem and a I return the image and a dictionary, for some reason it collates them. How do I stop this so it returns the same format as you provided when using the DataLoader class to iterate through?