I’m trying to use a sparse matrix in my model. The sparse matrix is fixed and not being updated. I’m using a custom function for this. However, I’m getting the following error when I use DataParallel. Any ideas how to solve this?
RuntimeError: zeros is not implemented for type torch.cuda.sparse.

And here is the code. It works if you comment out the DataParallel line.
import torch

class MyFunction(torch.autograd.Function):
  def forward(self, tensor, sparse_tensor):
    self.save_for_backward(sparse_tensor)
    print sparse_tensor.size(), tensor.size()
    return torch.matmul(sparse_tensor, tensor.t()).t()

  def backward(self, grad_output):
    sparse_tensor, = self.saved_tensors
    print sparse_tensor.size(), grad_output.size()
    return torch.matmul(sparse_tensor.t(), grad_output.t()), None


class MyCustomNet(torch.nn.Module):
  def __init__(self):
    super(MyCustomNet, self).__init__()
    i = torch.LongTensor([[0, 1, 1],
                          [2, 0, 2]])
    v = torch.FloatTensor([3, 4, 5])
    self.sparse_map = torch.nn.Parameter(torch.sparse.FloatTensor(i, v, torch.Size([2,3])))
    print torch.sparse.FloatTensor(i, v, torch.Size([2,3])).to_dense()
  def forward(self, x):
    self.func = MyFunction()
    return self.func(x, self.sparse_map)

net = MyCustomNet()
net.eval()
net.cuda()
net = torch.nn.DataParallel(net)
data = torch.autograd.Variable(torch.zeros((10, 3))).cuda()
target = torch.autograd.Variable(torch.zeros((10, 2))).cuda()


out = net(data)
loss_fn = torch.nn.L1Loss()
loss = loss_fn(out, target)
loss.backward()