If we are initialising a class instance of Module 1 in Module 2’s forward layer, then Module1’s parameters along with Layer 1 and Layer2’s don’t get registered.
class Layer1(nn.Module):
  def __init__(self, a, n, dim):
    super(Layer1, self).__init__()
    self.a = a
    self.n = n
    self.dim = dim
    self.w_input = nn.Parameter(torch.FloatTensor(self.dim, self.n,1,1))
    if a:
      self.wt = nn.Parameter(torch.FloatTensor(2 * self.dim, n, 1))

class Layer2(nn.Module):
  def __init__(self, n, dim):
    super(Layer1, self).__init__()
    self.n = n
    self.dim = dim
    self.g = nn.Parameter(torch.ones((1,dim,1,n)))


class Module1(nn.Module):
  def __init__(self, a, n, dim):
    super(Module1, self).__init__()
    self.a = a
    self.n = n
    self.dim = dim
    self.layer1 = Layer1(a, dim, dim)
    self.layer2 = Layer2(n, dim)

def forward(self, input):
    x = self.layer1(input)
    x = self.layer2(x)
    return x


 class Module2(nn.Module):
    def __init__():
       super(Module2, self).__init__()
       self.a = a
       self.n=n
       self.dim = dim
       self.layer3 = Layer3(a,n)
       
    def forward(self, input):
       x = self.layer3(input)
       inp_size = list(x.size())
       _, _, _, ndim = inp_size[0],inp_size[1],inp_size[2],inp_size[3]
       k = 2*(ndim-self.dim)
       if(k>1):
          layer_output = Module1(self.a, k, self.dim)
          output = layer_output(x)
       return output

The value of k (in Module2 forward layer) can only be determined from the values obtained in this forward layer, so we cannot define Module 1 instance in the init layer of Module 2.
So how do we register the parameters of the instance initialised in the forward layer?
Thanks!