I checked the above concept of initiating a trainable module inside the first forward pass only with the custom network. Please give feedback on its correctness.
from torch import nn
import torch


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


class network(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = convolution()
        self.conv2 = convolution()
        self.conv3 = convolution()

    def forward(self, x):
        x = self.conv1(x)
        x = self.conv2(x)
        x = self.conv3(x)
        return x


net = network()

print("*" * 20, "Before warm up", "*" * 20)
print("network:\n", net)

# Not going to print anything
for name, param in net.named_parameters():
    if param.requires_grad:
        print(name)

# warm up
rand_input_1 = torch.rand((16, 64, 7, 7))
a1 = net(rand_input_1)

print("*" * 20, "After warm up", "*" * 20)
print(net)
for name, param in net.named_parameters():
    if param.requires_grad:
        print(name, param.data.sum())

# random input with diffrent channel size will throw error
rand_input_2 = torch.rand((16, 64, 7, 7))
a2 = net(rand_input_2)

print("\nParameters is not initiated again in second call. It remains with same weight\n")
for name, param in net.named_parameters():
    if param.requires_grad:
        print(name, param.data.sum())

The output of the above snippet
******************** Before warm up ********************
network:
 network(
  (conv1): convolution()
  (conv2): convolution()
  (conv3): convolution()
)
******************** After warm up ********************
network(
  (conv1): convolution(
    (conv): Conv2d(64, 32, kernel_size=(3, 3), stride=(1, 1))
  )
  (conv2): convolution(
    (conv): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1))
  )
  (conv3): convolution(
    (conv): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1))
  )
)
conv1.conv.weight tensor(1.9705)
conv1.conv.bias tensor(0.1167)
conv2.conv.weight tensor(-6.0852)
conv2.conv.bias tensor(0.0791)
conv3.conv.weight tensor(-2.9391)
conv3.conv.bias tensor(0.1060)
Parameters is not initiated again in second call. It remains with same weight
conv1.conv.weight tensor(1.9705)
conv1.conv.bias tensor(0.1167)
conv2.conv.weight tensor(-6.0852)
conv2.conv.bias tensor(0.0791)
conv3.conv.weight tensor(-2.9391)
conv3.conv.bias tensor(0.1060)


This network can automatically configure itself after the first input data injection