From af0033c9e96a4a8dcc04f8fd737d0ccad3364f10 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Thu, 29 Nov 2018 11:59:29 +0100 Subject: [PATCH] updates --- train.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/train.py b/train.py index ec1b90a3..3bd778e3 100644 --- a/train.py +++ b/train.py @@ -58,8 +58,9 @@ def main(opt): model.load_state_dict(checkpoint['model']) if torch.cuda.device_count() > 1: - print('Using ', torch.cuda.device_count(), ' GPUs') - model = nn.DataParallel(model) + raise Exception('Multi-GPU not currently supported: https://github.com/ultralytics/yolov3/issues/21') + # print('Using ', torch.cuda.device_count(), ' GPUs') + # model = nn.DataParallel(model) model.to(device).train() # # Transfer learning (train only YOLO layers) @@ -87,8 +88,9 @@ def main(opt): load_weights(model, 'weights/darknet53.conv.74') if torch.cuda.device_count() > 1: - print('Using ', torch.cuda.device_count(), ' GPUs') - model = nn.DataParallel(model) + raise Exception('Multi-GPU not currently supported: https://github.com/ultralytics/yolov3/issues/21') + # print('Using ', torch.cuda.device_count(), ' GPUs') + # model = nn.DataParallel(model) model.to(device).train() # Set optimizer