This commit is contained in:
Glenn Jocher 2018-11-29 11:59:29 +01:00
parent d5331be0a0
commit af0033c9e9
1 changed files with 6 additions and 4 deletions

View File

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