From e5dc942feef32c6286a5d3d2169e389d74bfbcc4 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Mon, 4 Mar 2019 17:38:38 +0100 Subject: [PATCH] updates --- train.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/train.py b/train.py index 57b5d44b..099730f3 100644 --- a/train.py +++ b/train.py @@ -48,8 +48,8 @@ def train( # Load weights to resume from model.load_state_dict(checkpoint['model']) - if torch.cuda.device_count() > 1: - model = nn.DataParallel(model) + # if torch.cuda.device_count() > 1: + # model = nn.DataParallel(model) model.to(device).train() # Transfer learning (train only YOLO layers) @@ -75,8 +75,8 @@ def train( load_darknet_weights(model, weights + 'yolov3-tiny.conv.15') cutoff = 15 - if torch.cuda.device_count() > 1: - model = nn.DataParallel(model) + # if torch.cuda.device_count() > 1: + # model = nn.DataParallel(model) model.to(device).train() # Set optimizer