diff --git a/__pycache__/track_reid.cpython-39.pyc b/__pycache__/track_reid.cpython-39.pyc
index b351d96..8dad084 100644
Binary files a/__pycache__/track_reid.cpython-39.pyc and b/__pycache__/track_reid.cpython-39.pyc differ
diff --git a/contrast/feat_extract/config.py b/contrast/feat_extract/config.py
index f0cc387..2fc5c78 100644
--- a/contrast/feat_extract/config.py
+++ b/contrast/feat_extract/config.py
@@ -76,8 +76,8 @@ class Config:
     lr_decay = 0.95  # 0.98
     weight_decay = 5e-4
     loss = 'cross_entropy'  # ['focal_loss', 'cross_entropy']
-    device = torch.device('cuda:1' if torch.cuda.is_available() else 'cpu')
-    # device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
+    # device = torch.device('cuda:1' if torch.cuda.is_available() else 'cpu')
+    device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
 
     pin_memory = True  # if memory is large, set it True to speed up a bit
     num_workers = 4  # dataloader
diff --git a/models/__pycache__/experimental.cpython-39.pyc b/models/__pycache__/experimental.cpython-39.pyc
index 0abfeeb..babf431 100644
Binary files a/models/__pycache__/experimental.cpython-39.pyc and b/models/__pycache__/experimental.cpython-39.pyc differ
diff --git a/models/experimental.py b/models/experimental.py
index 11f75e2..a20cc3d 100644
--- a/models/experimental.py
+++ b/models/experimental.py
@@ -76,7 +76,7 @@ def attempt_load(weights, device=None, inplace=True, fuse=True):
 
     model = Ensemble()
     for w in weights if isinstance(weights, list) else [weights]:
-        ckpt = torch.load(attempt_download(w), map_location='cpu')  # load
+        ckpt = torch.load(attempt_download(w), map_location=device)  # load
         ckpt = (ckpt.get('ema') or ckpt['model']).to(device).float()  # FP32 model
 
         # Model compatibility updates