From cc42a206ef4c4d0cc4998d496123daa6953e0b15 Mon Sep 17 00:00:00 2001
From: Mohammad Khoshbin <khoshbin.mohammad.mk@gmail.com>
Date: Sat, 30 Jul 2022 20:13:18 +0430
Subject: [PATCH] fix cpu-only training CUDA error (#369)

---
 train.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/train.py b/train.py
index b87d687..0c5deaa 100644
--- a/train.py
+++ b/train.py
@@ -68,7 +68,7 @@ def train(hyp, opt, device, tb_writer=None):
     loggers = {'wandb': None}  # loggers dict
     if rank in [-1, 0]:
         opt.hyp = hyp  # add hyperparameters
-        run_id = torch.load(weights).get('wandb_id') if weights.endswith('.pt') and os.path.isfile(weights) else None
+        run_id = torch.load(weights, map_location=device).get('wandb_id') if weights.endswith('.pt') and os.path.isfile(weights) else None
         wandb_logger = WandbLogger(opt, Path(opt.save_dir).stem, run_id, data_dict)
         loggers['wandb'] = wandb_logger.wandb
         data_dict = wandb_logger.data_dict