Thanks, this is great! I’ve only recently learned about Lightning and hadn’t tried it yet, but this is very good motivation. (And also thanks for pointing out the validation shuffling - I was trying to replicate Keras’ settings, but it turns out that shuffle=True in Keras also only shuffles the training data.)
It turns out that the solution is very simple:
I initially thought that num_workers was unnecessary; because the dataset is small, I just transferred the whole thing to the GPU at the beginning (torch.Tensor(np.concatenate([class0_events, class1_events])[permutation]).to(device)). A loader with num_workers won’t even run in this case, because all the data is already loaded on the GPU.
However, it turns out to be faster not to do this, and to let the dataloader use multiple workers to load the relevant data to the GPU for each batch. I suspect this is because of shuffle=True; it’s better to use all the CPU cores to send just one batch’s data as contiguous register entries on the GPU. (As usual, it make sense in hindsight…) Lightning also takes care of the to(device) calls for the data automatically in the training loop, which is nice.