Process Pool freezes

1

The code below just freezes, and I don’t really know why. However, one of these small edits makes the code run just fine.

  • remove “weights=torchvision.models.DenseNet121_Weights.DEFAULT”
  • reduce the size of the returned tensor (i.e. 3, 224, 50 is a maximum)
from tqdm import tqdm
import torch
import torchvision
from torch.multiprocessing import Pool

def load_item(idx):
    print(idx)
    return torch.ones(3, 224, 224, dtype=torch.float32)
    
if __name__ == "__main__":
    model = torchvision.models.densenet121(weights=torchvision.models.DenseNet121_Weights.DEFAULT)
    with Pool(8) as p:
        dataset = p.map(load_item, tqdm(range(1000)))

The system I use

  • Ubuntu 20.04.4 LTS with 252 GB of RAM
  • python 3.8.18
  • torch 2.2.1+cu118
  • torchvision 0.17.1+cu118

I cannot replicate the bug on Windows.

Is there any particular reason why this code froze?