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