This question is about best practices for designing subclasses of nn.Module. I am trying to decrease unnecessary coupling in my code.
For example, say my network has a Conv2d layer and a Linear layer. The hyperparameters are: size of input image, number of conv. channels, kernel size, stride, number of FC layers. When building the model, I need to know the input size to the Linear layer in advance (as described in this answer 3). The arguments I must pass to initialize the Linear layer are coupled to other hyperparameters: whenever I change some hyperparameters of layers earlier in the chain(for example, image size or stride) I have to (or my code has to) recalculate the value. Presumably this is why there is demand for a Utility function for calculating the shape of a conv output 3. Conversely, the arguments needed to initialize a Conv2d layer are not coupled to hyperparameters earlier in the chain.
As a way of reducing coupling when writing custom modules, is it acceptable to move some initialization code from __init__() to forward()? forward() would then look like the pseudocode below.
def forward(self, x):
    if not self.initialized:
        # Initialize inner modules somehow using the input size
        self.weights = nn.Parameter(torch.randn(..., x.size(2), x.size(3), ...))
        self.initialized = True

    # Go on to usual forward() function
    ...

What would be the drawbacks of this approach? Are there other approaches to reducing coupling?