I am trying to move the image to the GPU during transforms.Compose but I am getting the error
transforms.Lambda(lambda x: x.to(device)),
RuntimeError: CUDA error: initialization error
I need to do this as I am using one of the trained models during Compose.transformation; something like:
device = torch.device('cuda')
transform = transforms.Compose([
transforms.Resize([256, 256]),
transforms.ToTensor(),
transforms.Lambda(lambda x: x.to(device)),
transforms.Lambda(lambda x: my_model(x.unsqueeze(0)).squeeze(0) )
])
Any ideas of how can I resolve this issue?
Thanks in advance