Well, reading this thread Deadlock with DataLoader and xarray/dask I found that adding this to your import
import dask
dask.config.set(scheduler='synchronous')
Solves the issue about stalling the Dataloader when num_workers > 0
I tested your code and confirmed that it stalls when num_workers > 0
. However, changing the dask config at the beginning of your script will work as it is intended.