Update utils.py

solve the multi-gpu training problem
2019-03-25 14:59:02 +08:00 · 2019-03-25 14:59:02 +08:00 · 4884508110
parent 386835d7ca
commit 4884508110
1 changed files with 1 additions and 1 deletions
--- a/utils/utils.py
+++ b/utils/utils.py
@ -285,7 +285,7 @@ def compute_loss(p, targets):  # predictions, targets

 def build_targets(model, targets, pred):
    # targets = [image, class, x, y, w, h]
-    if isinstance(model, nn.DataParallel):
+    if isinstance(model, nn.parallel.DistributedDataParallel):
        model = model.module
    yolo_layers = get_yolo_layers(model)