Hub models map_location=device (#3894)

* Hub models `map_location=device`

* cleanup

hubconf.py

@@ -36,13 +36,15 @@ def _create(name, pretrained=True, channels=3, classes=80, autoshape=True, verbo
 
     fname = Path(name).with_suffix('.pt')  # checkpoint filename
     try:
+        device = select_device(('0' if torch.cuda.is_available() else 'cpu') if device is None else device)
+
         if pretrained and channels == 3 and classes == 80:
-            model = attempt_load(fname, map_location=torch.device('cpu'))  # download/load FP32 model
+            model = attempt_load(fname, map_location=device)  # download/load FP32 model
         else:
             cfg = list((Path(__file__).parent / 'models').rglob(f'{name}.yaml'))[0]  # model.yaml path
             model = Model(cfg, channels, classes)  # create model
             if pretrained:
-                ckpt = torch.load(attempt_download(fname), map_location=torch.device('cpu'))  # load
+                ckpt = torch.load(attempt_download(fname), map_location=device)  # load
                 msd = model.state_dict()  # model state_dict
                 csd = ckpt['model'].float().state_dict()  # checkpoint state_dict as FP32
                 csd = {k: v for k, v in csd.items() if msd[k].shape == v.shape}  # filter
@@ -51,7 +53,6 @@ def _create(name, pretrained=True, channels=3, classes=80, autoshape=True, verbo
                     model.names = ckpt['model'].names  # set class names attribute
         if autoshape:
             model = model.autoshape()  # for file/URI/PIL/cv2/np inputs and NMS
-        device = select_device('0' if torch.cuda.is_available() else 'cpu') if device is None else torch.device(device)
         return model.to(device)
 
     except Exception as e: