CUDA/CPU error on Google Colab

Hi, I write a script based on pytorch that can transform a image to another one. It can work well on my pc, but since my GPU performance is too limited, I decide to run it on Google Colab. However, the same code cannot run on Colab.

Here is my code:

    # Use the cuda
    device = torch.device('cuda')    

    # Load Generator and send it to cuda
    G = UNet()
    G.cuda()
    G.load_state_dict(premodel['HT'])
    G.eval()
    print("loading success")
    G.to(device)
    # Get the images to do transform
    names = os.listdir(dataset_dir)
    names = names[:]
    for name in names:
        im_path = os.path.join(dataset_dir, name)

        # use PIL to open
        img = Image.open(im_path).convert('RGB')
        # print(img.size)

        # Then, convert this image to a tensor, then pass it to the device
        img_tensor = transforms.Compose([
                        transforms.ToTensor(),
                        transforms.Normalize(mean=[0.5, 0.5, 0.5],
                                             std=[0.5, 0.5, 0.5])
                    ])(img).to(device)
        #TO Make it from 3d to 4d
        img_tensor = img_tensor.view(1, img_tensor.shape[0], img_tensor.shape[1], img_tensor.shape[2])
        print(img_tensor.shape)
        # Next, pass this tensor to the transform net
        img_out = G(img_tensor)
        img_out = img_out.view(img_out.shape[1], img_out.shape[2], img_out.shape[3])
        img_out = transforms.Normalize(mean=[-1, -1, -1], std=[2, 2, 2])(img_out)


        # Pass the result to png format and save it to the result director
        img_png = transforms.ToPILImage()(img_out.detach().cpu()).convert('RGB')

        print("saving ", name)
        save_path = os.path.join(save_dir, name)
        img_png.save(save_path)
        img.close()
        img_png.close()
        del img_out
        del img_tensor

I think I have put all the model and tensors into the GPU, and it can run on my computer with a GPU. However, on Colab it always show the error:

Traceback (most recent call last):
File “demo_test.py”, line 85, in
img_out = transforms.Normalize(mean=[-1, -1, -1], std=[2, 2, 2])(img_out)
File “/usr/local/lib/python3.6/dist-packages/torchvision/transforms/transforms.py”, line 163, in call
return F.normalize(tensor, self.mean, self.std, self.inplace)
File “/usr/local/lib/python3.6/dist-packages/torchvision/transforms/functional.py”, line 208, in normalize
tensor.sub_(mean[:, None, None]).div_(std[:, None, None])
RuntimeError: expected backend CUDA and dtype Float but got backend CPU and dtype Float

So this runtime error comes from the line

img_out = transforms.Normalize(mean=[-1, -1, -1], std=[2, 2, 2])(img_out)

I think that error means img_out is not on the cuda, however I think I have sent it on the cuda. So can anyone tell me where does the problem come from? Thanks!!

Could you try to push the img_tensor to the GPU after the second normalization was performed?
Alternatively, try to create the mean and std as CUDA tensors.

I tried to push the second img_out to GPU again, but still gives the same error information; But it can run pretty well on my computer.

For create the mean and std as CUDA tensors, could you please give me some hint? I am pretty new to pytorch and have no idea how to do it. Thanks!

Sorry, I might have misunderstood your code, as it also runs fine of my machine.
Could you check the PyTorch and torchvision versions you are using in Colab?

Just for the sake of debugging, could you use this transform and see if it’s working?

img_out = transforms.Normalize(mean=torch.tensor([-1., -1., -1.], device='cuda'),
                               std=torch.tensor([2., 2., 2.], device='cuda'))(img_out)

I’m getting this error too.
I tried to pass everything it’s possible to ‘cuda’, but no better results…

Is your code also running fine locally and crashes on Google Colab?

I was trying in a computer without GPU and it was working fine.
After some searches I found that it looks like a known pytorch bug. See this github issue.
Then I checked all the variables I was trying to pass to GPU if they were in fact on GPU, and some of them were NOT!
Then I forced they to be on GPU since the instantiation and my code worked fine after this.
My problem was solved :smiley: