diff --git a/app/yolov5/train_server.py b/app/yolov5/train_server.py index a99990e..ab5ac4e 100644 --- a/app/yolov5/train_server.py +++ b/app/yolov5/train_server.py @@ -609,7 +609,7 @@ def main(opt,data_list,id,callbacks=Callbacks()): # DDP mode print('\n','opt.device:',opt.device,'\n') - device = opt.device #select_device(opt.device, batch_size=opt.batch_size) + device = torch.device(opt.device) #select_device(opt.device, batch_size=opt.batch_size) print('\n','device:',device,'\n') # if LOCAL_RANK != -1: # msg = 'is not compatible with YOLOv5 Multi-GPU DDP training'