We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
show_results with yolov5 + lightning training throws error
## 🐛 Bug
**Describe the bug**
`show_results` with yolo don't work if the model was training with pytorch lightning

**To Reproduce**
Train a yolo model with pytorch-lightning and try to call `show_results`

**Full stacktrace**
```python
---------------------------------------------------------------------------

RuntimeError                              Traceback (most recent call last)

<ipython-input-19-2cf4276b061d> in <module>()
----> 1 model_type.show_results(model, valid_ds, detection_threshold=.5)

9 frames

/usr/local/lib/python3.7/dist-packages/icevision/models/ultralytics/yolov5/show_results.py in show_results(model, dataset, detection_threshold, num_samples, ncols, denormalize_fn, show, device)
     41         show=show,
     42         detection_threshold=detection_threshold,
---> 43         device=device,
     44     )
     45 

/usr/local/lib/python3.7/dist-packages/icevision/models/base_show_results.py in base_show_results(predict_fn, model, dataset, num_samples, ncols, denormalize_fn, show, **predict_kwargs)
     19 ) -> None:
     20     records = random.choices(dataset, k=num_samples)
---> 21     preds = predict_fn(model, records, **predict_kwargs)
     22 
     23     show_preds(

/usr/local/lib/python3.7/dist-packages/icevision/models/ultralytics/yolov5/prediction.py in predict(model, dataset, detection_threshold, keep_images, device)
     48         detection_threshold=detection_threshold,
     49         keep_images=keep_images,
---> 50         device=device,
     51     )
     52 

/usr/local/lib/python3.7/dist-packages/torch/autograd/grad_mode.py in decorate_context(*args, **kwargs)
     24         def decorate_context(*args, **kwargs):
     25             with self.__class__():
---> 26                 return func(*args, **kwargs)
     27         return cast(F, decorate_context)
     28 

/usr/local/lib/python3.7/dist-packages/icevision/models/ultralytics/yolov5/prediction.py in _predict_batch(model, batch, records, detection_threshold, keep_images, device)
     24     model = model.eval().to(device)
     25 
---> 26     raw_preds = model(batch)[0]
     27     return convert_raw_predictions(
     28         batch=batch,

/usr/local/lib/python3.7/dist-packages/torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)
    725             result = self._slow_forward(*input, **kwargs)
    726         else:
--> 727             result = self.forward(*input, **kwargs)
    728         for hook in itertools.chain(
    729                 _global_forward_hooks.values(),

/usr/local/lib/python3.7/dist-packages/yolov5/models/yolo.py in forward(self, x, augment, profile)
    121             return torch.cat(y, 1), None  # augmented inference, train
    122         else:
--> 123             return self.forward_once(x, profile)  # single-scale inference, train
    124 
    125     def forward_once(self, x, profile=False):

/usr/local/lib/python3.7/dist-packages/yolov5/models/yolo.py in forward_once(self, x, profile)
    137                 print('%10.1f%10.0f%10.1fms %-40s' % (o, m.np, dt[-1], m.type))
    138 
--> 139             x = m(x)  # run
    140             y.append(x if m.i in self.save else None)  # save output
    141 

/usr/local/lib/python3.7/dist-packages/torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)
    725             result = self._slow_forward(*input, **kwargs)
    726         else:
--> 727             result = self.forward(*input, **kwargs)
    728         for hook in itertools.chain(
    729                 _global_forward_hooks.values(),

/usr/local/lib/python3.7/dist-packages/yolov5/models/yolo.py in forward(self, x)
     52 
     53                 y = x[i].sigmoid()
---> 54                 y[..., 0:2] = (y[..., 0:2] * 2. - 0.5 + self.grid[i]) * self.stride[i]  # xy
     55                 y[..., 2:4] = (y[..., 2:4] * 2) ** 2 * self.anchor_grid[i]  # wh
     56                 z.append(y.view(bs, -1, self.no))

RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!
```


--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `icevision/models/ultralytics/yolov5/prediction.py`
Content:
```
1 __all__ = ["predict", "predict_dl", "convert_raw_predictions"]
2 
3 from icevision.imports import *
4 from icevision.utils import *
5 from icevision.core import *
6 from icevision.data import *
7 from icevision.models.utils import _predict_dl
8 from icevision.models.ultralytics.yolov5.dataloaders import *
9 from yolov5.utils.general import non_max_suppression
10 
11 
12 @torch.no_grad()
13 def _predict_batch(
14     model: nn.Module,
15     batch: Sequence[torch.Tensor],
16     records: Sequence[BaseRecord],
17     detection_threshold: float = 0.25,
18     nms_iou_threshold: float = 0.45,
19     keep_images: bool = False,
20     device: Optional[torch.device] = None,
21 ) -> List[Prediction]:
22     device = device or model_device(model)
23 
24     batch = batch[0].to(device)
25     model = model.eval().to(device)
26 
27     raw_preds = model(batch)[0]
28     return convert_raw_predictions(
29         batch=batch,
30         raw_preds=raw_preds,
31         records=records,
32         detection_threshold=detection_threshold,
33         nms_iou_threshold=nms_iou_threshold,
34         keep_images=keep_images,
35     )
36 
37 
38 def predict(
39     model: nn.Module,
40     dataset: Dataset,
41     detection_threshold: float = 0.25,
42     nms_iou_threshold: float = 0.45,
43     keep_images: bool = False,
44     device: Optional[torch.device] = None,
45 ) -> List[Prediction]:
46     batch, records = build_infer_batch(dataset)
47     return _predict_batch(
48         model=model,
49         batch=batch,
50         records=records,
51         detection_threshold=detection_threshold,
52         nms_iou_threshold=nms_iou_threshold,
53         keep_images=keep_images,
54         device=device,
55     )
56 
57 
58 def predict_dl(
59     model: nn.Module,
60     infer_dl: DataLoader,
61     show_pbar: bool = True,
62     keep_images: bool = False,
63     **predict_kwargs,
64 ):
65     return _predict_dl(
66         predict_fn=_predict_batch,
67         model=model,
68         infer_dl=infer_dl,
69         show_pbar=show_pbar,
70         keep_images=keep_images,
71         **predict_kwargs,
72     )
73 
74 
75 def convert_raw_predictions(
76     batch,
77     raw_preds: torch.Tensor,
78     records: Sequence[BaseRecord],
79     detection_threshold: float,
80     nms_iou_threshold: float,
81     keep_images: bool = False,
82 ) -> List[Prediction]:
83     dets = non_max_suppression(
84         raw_preds, conf_thres=detection_threshold, iou_thres=nms_iou_threshold
85     )
86     dets = [d.detach().cpu().numpy() for d in dets]
87     preds = []
88     for det, record, tensor_image in zip(dets, records, batch):
89 
90         pred = BaseRecord(
91             (
92                 ScoresRecordComponent(),
93                 ImageRecordComponent(),
94                 InstancesLabelsRecordComponent(),
95                 BBoxesRecordComponent(),
96             )
97         )
98 
99         pred.detection.set_class_map(record.detection.class_map)
100         pred.detection.set_labels_by_id(det[:, 5].astype(int))
101         pred.detection.set_bboxes([BBox.from_xyxy(*xyxy) for xyxy in det[:, :4]])
102         pred.detection.set_scores(det[:, 4])
103 
104         if keep_images:
105             record.set_img(tensor_to_image(tensor_image))
106 
107         preds.append(Prediction(pred=pred, ground_truth=record))
108 
109     return preds
110
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/icevision/models/ultralytics/yolov5/prediction.py b/icevision/models/ultralytics/yolov5/prediction.py
--- a/icevision/models/ultralytics/yolov5/prediction.py
+++ b/icevision/models/ultralytics/yolov5/prediction.py
@@ -19,7 +19,17 @@
     keep_images: bool = False,
     device: Optional[torch.device] = None,
 ) -> List[Prediction]:
-    device = device or model_device(model)
+    # device issue addressed on discord: https://discord.com/channels/735877944085446747/770279401791160400/832361687855923250
+    if device is not None:
+        raise ValueError(
+            "For YOLOv5 device can only be specified during model creation, "
+            "for more info take a look at the discussion here: "
+            "https://discord.com/channels/735877944085446747/770279401791160400/832361687855923250"
+        )
+    grid = model.model[-1].grid[-1]
+    # if `grid.numel() == 1` it means the grid isn't initialized yet and we can't
+    # trust it's device (will always be CPU)
+    device = grid.device if grid.numel() > 1 else model_device(model)
 
     batch = batch[0].to(device)
     model = model.eval().to(device)

{"golden_diff": "diff --git a/icevision/models/ultralytics/yolov5/prediction.py b/icevision/models/ultralytics/yolov5/prediction.py\n--- a/icevision/models/ultralytics/yolov5/prediction.py\n+++ b/icevision/models/ultralytics/yolov5/prediction.py\n@@ -19,7 +19,17 @@\n     keep_images: bool = False,\n     device: Optional[torch.device] = None,\n ) -> List[Prediction]:\n-    device = device or model_device(model)\n+    # device issue addressed on discord: https://discord.com/channels/735877944085446747/770279401791160400/832361687855923250\n+    if device is not None:\n+        raise ValueError(\n+            \"For YOLOv5 device can only be specified during model creation, \"\n+            \"for more info take a look at the discussion here: \"\n+            \"https://discord.com/channels/735877944085446747/770279401791160400/832361687855923250\"\n+        )\n+    grid = model.model[-1].grid[-1]\n+    # if `grid.numel() == 1` it means the grid isn't initialized yet and we can't\n+    # trust it's device (will always be CPU)\n+    device = grid.device if grid.numel() > 1 else model_device(model)\n \n     batch = batch[0].to(device)\n     model = model.eval().to(device)\n", "issue": "show_results with yolov5 + lightning training throws error\n## \ud83d\udc1b Bug\r\n**Describe the bug**\r\n`show_results` with yolo don't work if the model was training with pytorch lightning\r\n\r\n**To Reproduce**\r\nTrain a yolo model with pytorch-lightning and try to call `show_results`\r\n\r\n**Full stacktrace**\r\n```python\r\n---------------------------------------------------------------------------\r\n\r\nRuntimeError                              Traceback (most recent call last)\r\n\r\n<ipython-input-19-2cf4276b061d> in <module>()\r\n----> 1 model_type.show_results(model, valid_ds, detection_threshold=.5)\r\n\r\n9 frames\r\n\r\n/usr/local/lib/python3.7/dist-packages/icevision/models/ultralytics/yolov5/show_results.py in show_results(model, dataset, detection_threshold, num_samples, ncols, denormalize_fn, show, device)\r\n     41         show=show,\r\n     42         detection_threshold=detection_threshold,\r\n---> 43         device=device,\r\n     44     )\r\n     45 \r\n\r\n/usr/local/lib/python3.7/dist-packages/icevision/models/base_show_results.py in base_show_results(predict_fn, model, dataset, num_samples, ncols, denormalize_fn, show, **predict_kwargs)\r\n     19 ) -> None:\r\n     20     records = random.choices(dataset, k=num_samples)\r\n---> 21     preds = predict_fn(model, records, **predict_kwargs)\r\n     22 \r\n     23     show_preds(\r\n\r\n/usr/local/lib/python3.7/dist-packages/icevision/models/ultralytics/yolov5/prediction.py in predict(model, dataset, detection_threshold, keep_images, device)\r\n     48         detection_threshold=detection_threshold,\r\n     49         keep_images=keep_images,\r\n---> 50         device=device,\r\n     51     )\r\n     52 \r\n\r\n/usr/local/lib/python3.7/dist-packages/torch/autograd/grad_mode.py in decorate_context(*args, **kwargs)\r\n     24         def decorate_context(*args, **kwargs):\r\n     25             with self.__class__():\r\n---> 26                 return func(*args, **kwargs)\r\n     27         return cast(F, decorate_context)\r\n     28 \r\n\r\n/usr/local/lib/python3.7/dist-packages/icevision/models/ultralytics/yolov5/prediction.py in _predict_batch(model, batch, records, detection_threshold, keep_images, device)\r\n     24     model = model.eval().to(device)\r\n     25 \r\n---> 26     raw_preds = model(batch)[0]\r\n     27     return convert_raw_predictions(\r\n     28         batch=batch,\r\n\r\n/usr/local/lib/python3.7/dist-packages/torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)\r\n    725             result = self._slow_forward(*input, **kwargs)\r\n    726         else:\r\n--> 727             result = self.forward(*input, **kwargs)\r\n    728         for hook in itertools.chain(\r\n    729                 _global_forward_hooks.values(),\r\n\r\n/usr/local/lib/python3.7/dist-packages/yolov5/models/yolo.py in forward(self, x, augment, profile)\r\n    121             return torch.cat(y, 1), None  # augmented inference, train\r\n    122         else:\r\n--> 123             return self.forward_once(x, profile)  # single-scale inference, train\r\n    124 \r\n    125     def forward_once(self, x, profile=False):\r\n\r\n/usr/local/lib/python3.7/dist-packages/yolov5/models/yolo.py in forward_once(self, x, profile)\r\n    137                 print('%10.1f%10.0f%10.1fms %-40s' % (o, m.np, dt[-1], m.type))\r\n    138 \r\n--> 139             x = m(x)  # run\r\n    140             y.append(x if m.i in self.save else None)  # save output\r\n    141 \r\n\r\n/usr/local/lib/python3.7/dist-packages/torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)\r\n    725             result = self._slow_forward(*input, **kwargs)\r\n    726         else:\r\n--> 727             result = self.forward(*input, **kwargs)\r\n    728         for hook in itertools.chain(\r\n    729                 _global_forward_hooks.values(),\r\n\r\n/usr/local/lib/python3.7/dist-packages/yolov5/models/yolo.py in forward(self, x)\r\n     52 \r\n     53                 y = x[i].sigmoid()\r\n---> 54                 y[..., 0:2] = (y[..., 0:2] * 2. - 0.5 + self.grid[i]) * self.stride[i]  # xy\r\n     55                 y[..., 2:4] = (y[..., 2:4] * 2) ** 2 * self.anchor_grid[i]  # wh\r\n     56                 z.append(y.view(bs, -1, self.no))\r\n\r\nRuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!\r\n```\r\n\n", "before_files": [{"content": "__all__ = [\"predict\", \"predict_dl\", \"convert_raw_predictions\"]\n\nfrom icevision.imports import *\nfrom icevision.utils import *\nfrom icevision.core import *\nfrom icevision.data import *\nfrom icevision.models.utils import _predict_dl\nfrom icevision.models.ultralytics.yolov5.dataloaders import *\nfrom yolov5.utils.general import non_max_suppression\n\n\n@torch.no_grad()\ndef _predict_batch(\n    model: nn.Module,\n    batch: Sequence[torch.Tensor],\n    records: Sequence[BaseRecord],\n    detection_threshold: float = 0.25,\n    nms_iou_threshold: float = 0.45,\n    keep_images: bool = False,\n    device: Optional[torch.device] = None,\n) -> List[Prediction]:\n    device = device or model_device(model)\n\n    batch = batch[0].to(device)\n    model = model.eval().to(device)\n\n    raw_preds = model(batch)[0]\n    return convert_raw_predictions(\n        batch=batch,\n        raw_preds=raw_preds,\n        records=records,\n        detection_threshold=detection_threshold,\n        nms_iou_threshold=nms_iou_threshold,\n        keep_images=keep_images,\n    )\n\n\ndef predict(\n    model: nn.Module,\n    dataset: Dataset,\n    detection_threshold: float = 0.25,\n    nms_iou_threshold: float = 0.45,\n    keep_images: bool = False,\n    device: Optional[torch.device] = None,\n) -> List[Prediction]:\n    batch, records = build_infer_batch(dataset)\n    return _predict_batch(\n        model=model,\n        batch=batch,\n        records=records,\n        detection_threshold=detection_threshold,\n        nms_iou_threshold=nms_iou_threshold,\n        keep_images=keep_images,\n        device=device,\n    )\n\n\ndef predict_dl(\n    model: nn.Module,\n    infer_dl: DataLoader,\n    show_pbar: bool = True,\n    keep_images: bool = False,\n    **predict_kwargs,\n):\n    return _predict_dl(\n        predict_fn=_predict_batch,\n        model=model,\n        infer_dl=infer_dl,\n        show_pbar=show_pbar,\n        keep_images=keep_images,\n        **predict_kwargs,\n    )\n\n\ndef convert_raw_predictions(\n    batch,\n    raw_preds: torch.Tensor,\n    records: Sequence[BaseRecord],\n    detection_threshold: float,\n    nms_iou_threshold: float,\n    keep_images: bool = False,\n) -> List[Prediction]:\n    dets = non_max_suppression(\n        raw_preds, conf_thres=detection_threshold, iou_thres=nms_iou_threshold\n    )\n    dets = [d.detach().cpu().numpy() for d in dets]\n    preds = []\n    for det, record, tensor_image in zip(dets, records, batch):\n\n        pred = BaseRecord(\n            (\n                ScoresRecordComponent(),\n                ImageRecordComponent(),\n                InstancesLabelsRecordComponent(),\n                BBoxesRecordComponent(),\n            )\n        )\n\n        pred.detection.set_class_map(record.detection.class_map)\n        pred.detection.set_labels_by_id(det[:, 5].astype(int))\n        pred.detection.set_bboxes([BBox.from_xyxy(*xyxy) for xyxy in det[:, :4]])\n        pred.detection.set_scores(det[:, 4])\n\n        if keep_images:\n            record.set_img(tensor_to_image(tensor_image))\n\n        preds.append(Prediction(pred=pred, ground_truth=record))\n\n    return preds\n", "path": "icevision/models/ultralytics/yolov5/prediction.py"}], "after_files": [{"content": "__all__ = [\"predict\", \"predict_dl\", \"convert_raw_predictions\"]\n\nfrom icevision.imports import *\nfrom icevision.utils import *\nfrom icevision.core import *\nfrom icevision.data import *\nfrom icevision.models.utils import _predict_dl\nfrom icevision.models.ultralytics.yolov5.dataloaders import *\nfrom yolov5.utils.general import non_max_suppression\n\n\n@torch.no_grad()\ndef _predict_batch(\n    model: nn.Module,\n    batch: Sequence[torch.Tensor],\n    records: Sequence[BaseRecord],\n    detection_threshold: float = 0.25,\n    nms_iou_threshold: float = 0.45,\n    keep_images: bool = False,\n    device: Optional[torch.device] = None,\n) -> List[Prediction]:\n    # device issue addressed on discord: https://discord.com/channels/735877944085446747/770279401791160400/832361687855923250\n    if device is not None:\n        raise ValueError(\n            \"For YOLOv5 device can only be specified during model creation, \"\n            \"for more info take a look at the discussion here: \"\n            \"https://discord.com/channels/735877944085446747/770279401791160400/832361687855923250\"\n        )\n    grid = model.model[-1].grid[-1]\n    # if `grid.numel() == 1` it means the grid isn't initialized yet and we can't\n    # trust it's device (will always be CPU)\n    device = grid.device if grid.numel() > 1 else model_device(model)\n\n    batch = batch[0].to(device)\n    model = model.eval().to(device)\n\n    raw_preds = model(batch)[0]\n    return convert_raw_predictions(\n        batch=batch,\n        raw_preds=raw_preds,\n        records=records,\n        detection_threshold=detection_threshold,\n        nms_iou_threshold=nms_iou_threshold,\n        keep_images=keep_images,\n    )\n\n\ndef predict(\n    model: nn.Module,\n    dataset: Dataset,\n    detection_threshold: float = 0.25,\n    nms_iou_threshold: float = 0.45,\n    keep_images: bool = False,\n    device: Optional[torch.device] = None,\n) -> List[Prediction]:\n    batch, records = build_infer_batch(dataset)\n    return _predict_batch(\n        model=model,\n        batch=batch,\n        records=records,\n        detection_threshold=detection_threshold,\n        nms_iou_threshold=nms_iou_threshold,\n        keep_images=keep_images,\n        device=device,\n    )\n\n\ndef predict_dl(\n    model: nn.Module,\n    infer_dl: DataLoader,\n    show_pbar: bool = True,\n    keep_images: bool = False,\n    **predict_kwargs,\n):\n    return _predict_dl(\n        predict_fn=_predict_batch,\n        model=model,\n        infer_dl=infer_dl,\n        show_pbar=show_pbar,\n        keep_images=keep_images,\n        **predict_kwargs,\n    )\n\n\ndef convert_raw_predictions(\n    batch,\n    raw_preds: torch.Tensor,\n    records: Sequence[BaseRecord],\n    detection_threshold: float,\n    nms_iou_threshold: float,\n    keep_images: bool = False,\n) -> List[Prediction]:\n    dets = non_max_suppression(\n        raw_preds, conf_thres=detection_threshold, iou_thres=nms_iou_threshold\n    )\n    dets = [d.detach().cpu().numpy() for d in dets]\n    preds = []\n    for det, record, tensor_image in zip(dets, records, batch):\n\n        pred = BaseRecord(\n            (\n                ScoresRecordComponent(),\n                ImageRecordComponent(),\n                InstancesLabelsRecordComponent(),\n                BBoxesRecordComponent(),\n            )\n        )\n\n        pred.detection.set_class_map(record.detection.class_map)\n        pred.detection.set_labels_by_id(det[:, 5].astype(int))\n        pred.detection.set_bboxes([BBox.from_xyxy(*xyxy) for xyxy in det[:, :4]])\n        pred.detection.set_scores(det[:, 4])\n\n        if keep_images:\n            record.set_img(tensor_to_image(tensor_image))\n\n        preds.append(Prediction(pred=pred, ground_truth=record))\n\n    return preds\n", "path": "icevision/models/ultralytics/yolov5/prediction.py"}]}