Merge pull request #143 from Lornatang/Fix-multi-GPU-training-errors

requirements.txt

@@ -2,7 +2,7 @@
 Cython
 numpy==1.17
 opencv-python
-torch>=1.5
+torch==1.4
 matplotlib
 pillow
 tensorboard

train.py

@@ -151,6 +151,7 @@ def train(hyp):
                                 world_size=1,  # number of nodes
                                 rank=0)  # node rank
         model = torch.nn.parallel.DistributedDataParallel(model)
+        # pip install torch==1.4.0+cu100 torchvision==0.5.0+cu100 -f https://download.pytorch.org/whl/torch_stable.html
 
     # Dataset
     dataset = LoadImagesAndLabels(train_path, imgsz, batch_size,