multi-fail.py

multi-fail.py

import torch
import numpy as np

class Model(torch.nn.Module):
    def forward(self, x):
        return torch.from_numpy(x).float().cuda()

ngpus = 2
x = np.ones((10, 10))

m = Model()
m.cuda()

out = torch.nn.parallel.data_parallel(m, x, range(ngpus))

print(x.shape) # (10, 10)
print(out.shape) # (20, 10)