Аргументы располагаются на разных графических процессорах при использовании nn.DataParallel (модель)

Question

Факел 0.4.1

Python 2.7.12

Я адаптировался Код NMP QC (с некоторыми проблемами совместимости) киспользовать несколько графических процессоров, поскольку мой графический процессор не может справиться с рабочей нагрузкой (сбой после завершения работы с VRAM)

Я новичок в pytorch, но я нашел учебное пособие по использованию nn.DataParallel (model) для реализации multi-gpu используйте

I , модифицированный main.py для использования nn.DataParallel (model). Области, которые я изменил, привязаны к "#NEW".

Код работает нормально даже в режиме нескольких графических процессоров, если работает на одном графическом процессоре, но получает ошибку «аргументы расположены на разных графических процессорах» при работе на двух или более графических процессорах

libibverbs: Warning: couldn't open config directory '/etc/libibverbs.d'.
libibverbs: Warning: no userspace device-specific driver found for /sys/class/infiniband_verbs/uverbs3
libibverbs: Warning: no userspace device-specific driver found for /sys/class/infiniband_verbs/uverbs2
libibverbs: Warning: no userspace device-specific driver found for /sys/class/infiniband_verbs/uverbs1
libibverbs: Warning: no userspace device-specific driver found for /sys/class/infiniband_verbs/uverbs0
Unexpected end of /proc/mounts line `overlay / overlay rw,relatime,lowerdir=/var/lib/docker/overlay2/l/QKSBQ5PAFDDC3OMBEELQQETALQ:/var/lib/docker/overlay2/l/WWYI3IDQPNXGON7AHODBPSTVXL:/var/lib/docker/overlay2/l/Q54I2HYS4TKH4LDJKBTVTGWWO6:/var/lib/docker/overlay2/l/IUV2LFPNMPOS3MREOTT52TKL54:/var/lib/docker/overlay2/l/DB5GBUCI3DCBPX6TJG3O337YVB:/var/lib/docker/overlay2/l/DNYKXCZJH5FMFNJLNGYJJ2ITPI:/var/lib/docker/overlay2/l/7DZCTDVNSTPJISGW65UG7U3F75:/var/lib/docker/overlay2/l/VOEQO652VS63NLDLZZ4TCIJLO6:/var/lib/docker/overlay2/l/4SI6ZCRUIORG5'
Traceback (most recent call last):
  File "main.py", line 332, in <module>
    main()
  File "main.py", line 190, in main
    train(train_loader, model, criterion, optimizer, epoch, evaluation, logger)
  File "main.py", line 251, in train
    output = model(g, h, e)
  File "/usr/local/lib/python2.7/dist-packages/torch/nn/modules/module.py", line 477, in __call__
    result = self.forward(*input, **kwargs)
  File "/usr/local/lib/python2.7/dist-packages/torch/nn/parallel/data_parallel.py", line 123, in forward
    outputs = self.parallel_apply(replicas, inputs, kwargs)
  File "/usr/local/lib/python2.7/dist-packages/torch/nn/parallel/data_parallel.py", line 133, in parallel_apply
    return parallel_apply(replicas, inputs, kwargs, self.device_ids[:len(replicas)])
  File "/usr/local/lib/python2.7/dist-packages/torch/nn/parallel/parallel_apply.py", line 77, in parallel_apply
    raise output
RuntimeError: arguments are located on different GPUs at /pytorch/aten/src/THC/generic/THCTensorMathBlas.cu:236

Поскольку я былотправляя входные данные по одному, а не сразу, как в учебнике, я проверил с помощью .get_device (), который подтвердил, что все 4 передаваемых аргумента (g, h, e, target) были на одном устройстве (устройство 0)