Fix DDP bug in single process multiple device use cases

2020-06-20 13:00:03 +08:00 · 2020-06-20 13:00:03 +08:00 · 899f1d4bde
parent 0a09882e4f
commit 899f1d4bde
2 changed files with 2 additions and 1 deletions
--- 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
--- 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,