out = dp.data_parallel(l, i, (0, 1))
self.assertEqual(out, l(i))
+ @unittest.skipIf(not TEST_MULTIGPU, "multi-GPU not supported")
+ @skipIfRocm
+ def test_data_parallel_model_device(self):
+ l = nn.Linear(2, 2)
+ error_msg = "module must have its parameters and buffers on device %d "
+ self.assertRaisesRegex(
+ RuntimeError, error_msg % (0), lambda: nn.DataParallel(l))
+ self.assertRaisesRegex(
+ RuntimeError, error_msg % (0), lambda: nn.DataParallel(l.cuda(1)))
+ self.assertRaisesRegex(
+ RuntimeError, error_msg % (1),
+ lambda: nn.DataParallel(l.cuda(), device_ids=[1, 0]))
+
+ nn.DataParallel(l.cuda())
+ nn.DataParallel(l.cuda(1), device_ids=[1, 0])
+
+ s = nn.Sequential(l.cpu())
+ self.assertRaisesRegex(
+ RuntimeError, error_msg % (0), lambda: nn.DataParallel(s))
+
+ s = nn.Sequential(deepcopy(l), l.cuda())
+ self.assertRaisesRegex(
+ RuntimeError, error_msg % (0), lambda: nn.DataParallel(s))
+
+ s = nn.Sequential(l.cuda(), deepcopy(l).cuda(1))
+ self.assertRaisesRegex(
+ RuntimeError, error_msg % (0), lambda: nn.DataParallel(s))
+ self.assertRaisesRegex(
+ RuntimeError, error_msg % (1),
+ lambda: nn.DataParallel(s, device_ids=[1, 0]))
+
+ s = nn.Sequential(l.cuda(), deepcopy(l).cuda())
+ nn.DataParallel(s)
+
+ s = nn.Sequential(l.cuda(1), deepcopy(l).cuda(1))
+ nn.DataParallel(s, device_ids=[1, 0])
+
@unittest.skipIf(not TEST_MULTIGPU or not PY3, "multi-GPU not supported")
@skipIfRocm
def test_data_parallel_model_no_refcycles(self):
import operator
import torch
import warnings
+from itertools import chain
from ..modules import Module
from .scatter_gather import scatter_kwargs, gather
from .replicate import replicate
Example::
>>> net = torch.nn.DataParallel(model, device_ids=[0, 1, 2])
- >>> output = net(input_var)
+ >>> output = net(input_var) # input_var can be on any device, including CPU
"""
# TODO: update notes/cuda.rst when this class handles 8+ GPUs well
if output_device is None:
output_device = device_ids[0]
+ if not all(t.is_cuda and t.device.index == device_ids[0]
+ for t in chain(module.parameters(), module.buffers())):
+ raise RuntimeError("module must have its parameters and buffers "
+ "on device %d (device_ids[0])" % device_ids[0])
+
self.dim = dim
self.module = module
self.device_ids = list(map(lambda x: _get_device_index(x, True), device_ids))