From 35396adc9cf3908ac65e0b146ba8e74f7c6fdbaa Mon Sep 17 00:00:00 2001
From: perry0418 <34980036+perry0418@users.noreply.github.com>
Date: Thu, 21 Mar 2019 12:02:57 +0800
Subject: [PATCH] Update train.py

solve the Multi-GPU --resume Error #138
https://github.com/ultralytics/yolov3/issues/138
---
 train.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/train.py b/train.py
index b791d7e4..7ec5a46c 100644
--- a/train.py
+++ b/train.py
@@ -32,7 +32,7 @@ def train(
     train_path = parse_data_cfg(data_cfg)['train']
 
     # Initialize model
-    model = Darknet(cfg, img_size)
+    model = Darknet(cfg, img_size).to(device)
 
     # Get dataloader
     dataloader = LoadImagesAndLabels(train_path, batch_size, img_size, augment=True)
@@ -43,7 +43,7 @@ def train(
     start_epoch = 0
     best_loss = float('inf')
     if resume:
-        checkpoint = torch.load(latest, map_location='cpu')
+        checkpoint = torch.load(latest, map_location=device)
 
         # Load weights to resume from
         model.load_state_dict(checkpoint['model'])