9a8d91022cfde0a380d8700be76d6ddbd716a2af,robosat/tools/train.py,,train,#,116

Before Change

        for mask, output in zip(masks, outputs):
            mask = mask.data.cpu().numpy()
            prediction = output.data.max(0)[1].cpu().numpy()
            iou.add(mask.ravel(), prediction.ravel())

    assert num_samples > 0, "dataset contains training images and labels"

After Change

        for mask, output in zip(masks, outputs):
            prediction = output.detach()
            prediction.requires_grad = False
            iou.add(mask.float(), prediction.max(0)[1].float())

    assert num_samples > 0, "dataset contains training images and labels"

    return {"loss": running_loss / num_samples, "iou": iou.get()}

Instances