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