672b2668ce155a0f7c1a56f3e46a3fc7205f3cb3,models/segmentation/unet.py,,,#,108

Before Change

        model.cuda()

    images = torch.randn(4, 3, 256, 256)
    if torch.cuda.is_available():
        images = images.cuda()

    out = model.forward(images)
    print(out.size())

After Change

if __name__ == "__main__":
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    model = UNet(num_classes=1).to(device)
    images = torch.randn(4, 3, 256, 256).to(device)

Instances