I’m currently working on porting code from Keras to PyTorch. I’m working with many GPUs and CPUs so it’s important to have batch generation happening in parallel.  My problem is that I’m trying to use the num_workers argument on the DataLoader class, but am meeting with errors.  Current relevant toy code:
import torch
torch.multiprocessing.set_start_method('spawn')
from torch.utils.data import Dataset, DataLoader

X_train = torch.randn(100,100)
y_train = torch.randn(100,100)
w_train = torch.randn(100,100)

class My_Dataset(Dataset):
     def __init__(self, x_input, y_labels, w_labels):
        self.y_labels = y_labels
        self.w_labels = w_labels
        self.x_input = x_input
        self.len = self.x_input.shape[0]

    def __len__(self):
        return self.len

    def __getitem__(self, idx):
        if __name__ == '__main__':
            print('Error! being run by {}'.format(__name__))
        
        X = self.x_input[idx]
        y = self.y_labels[idx]
        w = self.w_labels[idx]
        return X, y, w

training_dataset = My_Dataset(X_train, y_train, w_train)
training_dataloader = DataLoader(training_dataset, batch_size=10, shuffle=True, num_workers=15, pin_memory=False)

for batch_num, (inputs, labels_y, labels_w) in enumerate(training_dataloader):
    (inputs, labels_y, labels_w) =(inputs.cuda(), labels_y.cuda(), labels_w.cuda()) 
    print(inputs)

Removing torch.multiprocessing.set_start_method(‘spawn’) causes the code to run, but the batch generation runs in the main process (the error message I wrote into the dataset prints, also on my non-toy problem it takes unacceptably long).  I’ve tried both ‘spawn’ and ‘forkserver’, but both fail and I can’t even figure out how to get a relevant error message.  Right now the error just says one of the processes failed and setting num_workers =0 might get a better message.
Python version = 3.7.9, torch version = 1.7.0