Ошибка при попытке запустить Pytorch на нескольких графических процессорах

Question

Я пытаюсь создать сценарий с довольно простой очередью производителей/потребителей. Я использую это в системе с двумя графическими процессорами A4000. Ниже приведен соответствующий код.

      import torch
from torch.multiprocessing import Process, set_start_method, Queue

def main():
    input_data_queue = Queue(25)
    send_data_queue = Queue(5)

    for i in range(torch.cuda.device_count()):
        Process_Data(input_data_queue, send_data_queue, i)

    ....

class Process_Data:
    def __init__(self, in_q, out_q, gpu_id):
        self.in_queue = in_q
        self.out_queue = out_q
        self.gpu_id = gpu_id

        self.model = torch.hub.load('ultralytics/yolov5', 'custom', path='best.pt').to(torch.device(self.gpu_id))
        self.model.eval()
        ....

if __name__ == "__main__":
    set_start_method('spawn')
    main()

Я всегда получаю ошибку:

      Traceback (most recent call last):
File "<string>", line 1, in <module>
File "/usr/lib/python3.8/multiprocessing/spawn.py", line 116, in spawn_main
  exitcode = _main(fd, parent_sentinel)
File "/usr/lib/python3.8/multiprocessing/spawn.py", line 126, in _main
  self = reduction.pickle.load(from_parent)
File "/usr/local/lib/python3.8/dist-packages/torch/multiprocessing/reductions.py", line 111, in rebuild_cuda_tensor
  storage = storage_cls._new_shared_cuda(
File "/usr/local/lib/python3.8/dist-packages/torch/storage.py", line 630, in _new_shared_cuda
  return eval(cls.__module__)._UntypedStorage._new_shared_cuda(*args, **kwargs)
RuntimeError: CUDA error: invalid device ordinal
CUDA kernel errors might be asynchronously reported at some other API call,so the stacktrace below might be incorrect.
For debugging consider passing CUDA_LAUNCH_BLOCKING=1.

Для вызова устройства я пробовал:

1. Сначала создайте модель с помощью torch.device(0), затем еще один класс с помощью torch.device(1)
2. torch.device("cuda:0") затем torch.device("cuda:1")
3. torch.device("cuda") затем torch.device("cuda")
4. torch.device("cuda", 0), затем torch.device("cuda", 1)

Все варианты, которые я могу найти в документации, дают одну и ту же ошибку.

Как я могу получить две модели, работающие на двух графических процессорах, с общей рабочей очередью?