It is working.
Furthermore, I got this quirky idea of creating a convolution module where we don’t need to specify the input_channel explicitly and initiating module inside first forward pass.
class convolution(nn.Module):

    def __init__(self):
         super().__init__()
         self.initialized = True

    def forward(self,x):
         if self.initialized:
              in_channel = x.shape[1]
              self.conv = nn.Conv2d(in_channel, 32, 3)
              self.initialized = False
        
         x = self.conv(x)

         return x


Will it work as intended?