Merge pull request #145 from perry0418/master

Update train.py
2019-03-21 12:04:50 +02:00
parent d661fba8ae 35396adc9c
commit ad49e70f47
1 changed files with 2 additions and 2 deletions
@@ -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'])