Solved by ptrblck in post #2 
              
              
                You could use persistent_workers=True in the DataLoader creation and iterate it directly via: 
for batch in loader:
    ...

which will make sure to keep the workers alive.