PyTorch: How to create a Parameter without specifying the dimension
Say I want to defined a module. In this module, the __init__() function will create a Parameter called self.weight without known the input_dim of the module. My question is, how can I expand the self.weight and initialize it when I first call the forward() function?
For example, I want my module looks like this:
class MyModel(torch.nn.Module):
    def __init__(self, out_dim):
        super(MyModel, self).__init__()
        # I don't know the input_dim yet
        self.weight = torch.nn.Parameter(torch.FloatTensor(None, out_dim))
        self.init_weight = False
    
    def init_parameters(self, in_dim):
        # what should I do in this function? 
        # Is this correct?
        self.weight = self.weight.expand(in_dim, -1)
        torch.nn.init.xvaier_normal_(self.weight)
        self.init_weight = True

    def forward(self, X):
        if not self.init_weight:
            # first call, so now I can initialize the weight since I know the input_dim 
            self.init_parameters(X.shape[1])
        
        # do some forward ops
        return torch.sigmoid(torch.matmul(X, self.weight))        

And my training code looks like this (The parameter self.weight is passed to the optimizer after I create the model):
def train(X_train, y_train):
    model = MyModel(y_train.shape[1])
    optimize = torch.optim.Adam(model.parameters(), lr=1e-3)
    loss_fn = torch.nn.MSELoss()

    for epoch in range(10000):
        optimize.zero_grad()

        prediction = model(X_train)
        loss = loss_fn(prediction, y_train)
        loss.backward()

        optimize.step()


After all, it works for me using the way I explained in the comments - to allocate the weights parameter right in the init_parameters function.
import torch

class MyModel(torch.nn.Module):
    def __init__(self, out_dim):
        super(MyModel, self).__init__()
        self.weight = torch.nn.Parameter(torch.FloatTensor([0.0]))
        self.out_dim = out_dim
        self.init_weight = False
    
    def init_parameters(self, in_dim):
        self.weight = torch.nn.Parameter(torch.FloatTensor(in_dim, self.out_dim), requires_grad=True)
        torch.nn.init.xavier_normal_(self.weight)
        self.init_weight = True

    def forward(self, X):
        if not self.init_weight:
            # first call, so now I can initialize the weight since I know the input_dim 
            self.init_parameters(X.shape[1])
        
        # do some forward ops
        result = torch.sigmoid(torch.matmul(X, self.weight)) 
        print(X.shape, result.shape)
        return result
    
def train(X_train, y_train):
    model = torch.nn.Sequential(MyModel(out_dim=100), MyModel(out_dim=20))
    optimize = torch.optim.Adam(model.parameters(), lr=1e-3)
    loss_fn = torch.nn.MSELoss()

    for epoch in range(10000):
        #print('.', end='')
        optimize.zero_grad()

        prediction = model(X_train)
        loss = loss_fn(prediction, y_train)
        loss.backward()

        optimize.step()
        
batch_size, in_dim, out_dim = 100, 5, 20
X_train=torch.randn((batch_size, in_dim))
y_train=torch.randn((batch_size, out_dim))
train(X_train, y_train)

