To(device) gives an error when used inside transforms.Compose

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

If you are using multiple workers in your DataLoader, multiple initializations of the CUDA context might yield this error.
Try to use the spawn method for multiprocessing as described here or alternatively use a single worker (num_workers=0).

1 Like