Commit
·
899f1d4
1
Parent(s):
0a09882
Fix DDP bug in single process multiple device use cases
Browse files- requirements.txt +1 -1
- train.py +1 -0
requirements.txt
CHANGED
@@ -2,7 +2,7 @@
|
|
2 |
Cython
|
3 |
numpy==1.17
|
4 |
opencv-python
|
5 |
-
torch
|
6 |
matplotlib
|
7 |
pillow
|
8 |
tensorboard
|
|
|
2 |
Cython
|
3 |
numpy==1.17
|
4 |
opencv-python
|
5 |
+
torch==1.4
|
6 |
matplotlib
|
7 |
pillow
|
8 |
tensorboard
|
train.py
CHANGED
@@ -151,6 +151,7 @@ def train(hyp):
|
|
151 |
world_size=1, # number of nodes
|
152 |
rank=0) # node rank
|
153 |
model = torch.nn.parallel.DistributedDataParallel(model)
|
|
|
154 |
|
155 |
# Dataset
|
156 |
dataset = LoadImagesAndLabels(train_path, imgsz, batch_size,
|
|
|
151 |
world_size=1, # number of nodes
|
152 |
rank=0) # node rank
|
153 |
model = torch.nn.parallel.DistributedDataParallel(model)
|
154 |
+
# pip install torch==1.4.0+cu100 torchvision==0.5.0+cu100 -f https://download.pytorch.org/whl/torch_stable.html
|
155 |
|
156 |
# Dataset
|
157 |
dataset = LoadImagesAndLabels(train_path, imgsz, batch_size,
|