diff --git a/utils/utils.py b/utils/utils.py index 1f4e03b0..76b25df2 100755 --- 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)