Multiprocess dataloader using TensorDataset

Hello! I use the following code to load a tensordataset from a pkl file. However, when I set the data workers in DataLoader to 2, I will encounter a GPU error
“[W CUDAGuardImpl.h:62] Warning: CUDA warning: invalid device ordinal (function uncheckedSetDevice)
[W CUDAGuardImpl.h:46] Warning: CUDA warning: driver shutting down (function uncheckedGetDevice)”
How to solve this problem?

   train_image=torch. Load(train_image_dir)
   train_label=torch.load(train_label_dir)
   test_image=torch.load(test_image_dir)
   test_label = torch.load(test_label_dir)

   train_ds = TensorDataset(train_image, train_label)
   test_ds = TensorDataset(test_image, test_label)

   train_ds_len = len(train_ds)
   test_ds_len = len(test_ds)

   print(train_ds_len)
   print(test_ds_len)

   # for train
   train_dl = DataLoader(dataset=train_ds, batch_size=batch_size,
                         shuffle=True, num_workers=data_workers, drop_last=True)
   # for evaluate
   test_dl = DataLoader(dataset=test_ds, batch_size=batch_size,
                        shuffle=False, num_workers=data_workers, drop_last=False)

Based on the described behavior I would guess your Dataset is trying to push tensors directly to a GPU and is using an invalid device index. Could you check the stacktrace and narrow down where this operation is performed and post the corresponding code here, please?