Here’s another way that seems to work. Like your example with register_parameter, a forward pass is required before initializing the optimizer. Do you see any clear advantages or disadvantages?
class DynamicLinear(nn.Module):
    def __init__(self, output_dim):
        super(DynamicLinear, self).__init__()
        self.output_dim = output_dim

    def forward(self, inputs):
        if not hasattr(self, '_linear'):
            input_dim = inputs.shape[-1]
            self._linear = nn.Linear(input_dim, self.output_dim)
        return self._linear(inputs)