glenn-jocher commited on
Commit
f2d97eb
1 Parent(s): 3749573

Remove DDP MultiHeadAttention fix (#3768)

Browse files
Files changed (1) hide show
  1. train.py +1 -3
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