Yes I am using nn.DataParallel , sorry for sharing small snippets, it is a fairly large script and my research.
self.W_theta = nn.Parameter(torch.rand(1, 1, 1,  in_channels, out_channels, 1, 1, 1))
self.W_hat = nn.Parameter(torch.rand(1, 1, 1,  in_channels, out_channels, 3, 1, 1))
Input is a custom dataset with 2 channels. Everything is fine except this section I have done many experiments before adding this part.