mirrors
/
yolov7
mirror of https://github.com/WongKinYiu/yolov7.git
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

fix cpu-only training CUDA error (#369)

pull/355/merge
Mohammad Khoshbin 2022-07-30 20:13:18 +04:30 committed by GitHub
parent 6bacefff5c
commit cc42a206ef
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 1 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
train.py
View File

@ -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