@ptrblck
Actually for my conv2d function I am using autograd Functions. Like below
class Conv2d_function(Function):
@staticmethod
def forward(ctx, x1,x2):
    
    y1 = torch.tensor([1,2,3]).to(device)
    y2 = torch.tensor([4,5,6]).to(device)
    ....

class Conv2d(nn.Module):
def init(self, in_channels …):
…
…
super(Conv2d, self).init()
def forward(self, input):
    return Conv2d_function.apply(input)

Creating model using custom conv2d class
class Model():
layer1 = Conv2d(…)
layer2 = Conv2d(…)
model1 = Model()
So when I use Dataparallel for this model1 I get this error “RuntimeError: Expected tensor for argument #1 ‘input’ to have the same device as tensor for argument #2 ‘weight’; but device 0 does not equal 1 (while checking arguments for cudnn_batch_norm)”
Is this error occurring because in the forward function of Conv2d_function class I am defining a tensor using .to(device)?
Actually the tensor y1 and y2 depend on my input to the forward function of class Conv2d so I can’t define those tensor in the init of Conv2d class as register_buffer or Parameter. So I can only define those in my forward function of Conv2d_function class. But with data-parallel I get this above error which I feel is because I am transferring these tensors to GPU in my forward pass.