Dynamic parameter declaration in forward function

2

You can create parameters in the forward function too. Just guard them with an if to prevent reassigning at every iteration:

class MyModule(nn.Module):
    def __init__(self):
        # you need to register the parameter names earlier
        self.register_parameter('weight', None)

    def forward(self, input):
        if self.weight is None:
            self.weight = nn.Parameter(torch.randn(input.size()))
        return self.weight @ input
10 Likes