From 899f1d4bde374b0d98fe56c767cf1985d21cb72f Mon Sep 17 00:00:00 2001
From: Lornatang <liuchangyu1111@gmail.com>
Date: Sat, 20 Jun 2020 13:00:03 +0800
Subject: [PATCH] Fix DDP bug in single process multiple device use cases

---
 requirements.txt | 2 +-
 train.py         | 1 +
 2 files changed, 2 insertions(+), 1 deletion(-)

diff --git a/requirements.txt b/requirements.txt
index a0130d2..52846b8 100755
--- a/requirements.txt
+++ b/requirements.txt
@@ -2,7 +2,7 @@
 Cython
 numpy==1.17
 opencv-python
-torch>=1.5
+torch==1.4
 matplotlib
 pillow
 tensorboard
diff --git a/train.py b/train.py
index f545ae1..94ccf57 100644
--- a/train.py
+++ b/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,