I have a Dataset object that yields data generated from a pre-trained model that I load at the initialization of the dataset object. The model generating these data is placed on a separate GPU from the main model that will be optimized. Since the generation process is on a separate GPU, I would like to be able to generate more examples in parallel while the main model is performing an optimization step. I tried setting num_workers of the DataLoader object to a value greater than zero but I get the following error:
  File "/usr/local/lib/python3.8/dist-packages/torch/utils/data/dataloader.py", line 279, in __iter__
    return _MultiProcessingDataLoaderIter(self)
  File "/usr/local/lib/python3.8/dist-packages/torch/utils/data/dataloader.py", line 719, in __init__
    w.start()
  File "/usr/lib/python3.8/multiprocessing/process.py", line 121, in start
    self._popen = self._Popen(self)
  File "/usr/lib/python3.8/multiprocessing/context.py", line 224, in _Popen
    return _default_context.get_context().Process._Popen(process_obj)
  File "/usr/lib/python3.8/multiprocessing/context.py", line 284, in _Popen
    return Popen(process_obj)
  File "/usr/lib/python3.8/multiprocessing/popen_spawn_posix.py", line 32, in __init__
    super().__init__(process_obj)
  File "/usr/lib/python3.8/multiprocessing/popen_fork.py", line 19, in __init__
    self._launch(process_obj)
  File "/usr/lib/python3.8/multiprocessing/popen_spawn_posix.py", line 47, in _launch
    reduction.dump(process_obj, fp)
  File "/usr/lib/python3.8/multiprocessing/reduction.py", line 60, in dump
    ForkingPickler(file, protocol).dump(obj)
TypeError: cannot pickle 'module' object

Any suggestions on how to make this process faster in general?