I encountered a strange issue and the simplified version of my codes are as below:
import torch
import torch.nn as nn
import math

class MyLinear(nn.Module):
  def __init__(self, nin, nout):
    super(MyLinear, self).__init__()
    self.nout = nout
    self.nin = nin
    self.weight = nn.Parameter(torch.randn(self.nout, self.nin))
    self.reset_parameters()

  def reset_parameters(self):
    stdv = 1. / math.sqrt(self.weight.size(1))
    self.weight.data.uniform_(-stdv, stdv)
  
  def forward(self, x):
    # my_regularization = torch.abs(self.weight).mean().reshape(1)
    my_regularization = torch.abs(self.weight).mean()
    return torch.nn.functional.linear(x, self.weight), my_regularization

model = MyLinear(10, 1).cuda()
model = nn.DataParallel(model)
criterion = nn.MSELoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.01, weight_decay=0.01, momentum=0.1)

for i in range(100):
  data = torch.randn(100, 10).cuda()
  target = torch.randn(100, 1).cuda()
  output, my_regularization = model(data)
  print(output.shape, my_regularization.shape)
  loss = criterion(output, target)
  loss = loss + my_regularization
  optimizer.zero_grad()
  loss.backward()
  optimizer.step()

I’ve implemented a special version of Linear layer with customized regularization term.
And my Linear layer return both the output and the regularization term which will be optimized as part of the loss function.
When I employ the nn.DataParallel to perform multi-GPUs training, the program gives an error:
RuntimeError: dimension specified as 0 but tensor has no dimensions

The error indicates that my_regularization term has no axis.
But when I reshape it from an scalar to vector with .reshape(1), another error occurs:
RuntimeError: grad can be implicitly created only for scalar outputs

Because my_regularization strangely has shape=[2], which I think should be [1].
I have 2 GPUs on board and my_regularization may be the concatenation of both.
But why don’t the shape of output change from [100, 1] to [200, 1] as well?
In my oppinion they are both output of MyLinear layer, so what causes their different behaviours?
Could you folks give me any hints about solving this issue?