glenn-jocher
commited on
Commit
•
f2d97eb
1
Parent(s):
3749573
Remove DDP MultiHeadAttention fix (#3768)
Browse files
train.py
CHANGED
@@ -252,9 +252,7 @@ def train(hyp, # path/to/hyp.yaml or hyp dictionary
|
|
252 |
|
253 |
# DDP mode
|
254 |
if cuda and RANK != -1:
|
255 |
-
model = DDP(model, device_ids=[LOCAL_RANK], output_device=LOCAL_RANK
|
256 |
-
# nn.MultiheadAttention incompatibility with DDP https://github.com/pytorch/pytorch/issues/26698
|
257 |
-
find_unused_parameters=any(isinstance(layer, nn.MultiheadAttention) for layer in model.modules()))
|
258 |
|
259 |
# Model parameters
|
260 |
hyp['box'] *= 3. / nl # scale to layers
|
|
|
252 |
|
253 |
# DDP mode
|
254 |
if cuda and RANK != -1:
|
255 |
+
model = DDP(model, device_ids=[LOCAL_RANK], output_device=LOCAL_RANK)
|
|
|
|
|
256 |
|
257 |
# Model parameters
|
258 |
hyp['box'] *= 3. / nl # scale to layers
|