Pytorch - What module should I use to multiply the output of a layer using Sequential?

Question

While defining a neural network using nn.Module, in the forward function I can multiply the output of the final layer using:

def forward(self, x):
    ...
    x = torch.mul(x, self.max_action)
    return x

I am trying to do the same but instead using nn.Sequential method to define the neural network

model = nn.Sequential()
model.add_module(...
...
model.add_module(name='activation_output', module=?)

What should I use there to have the previous layer multiply by the scalar self.max_action ? Or should I build the sequential model in a different way ?

Answer 1

You could define a custom nn.Module layer:

class Multiply(nn.Module):
    def __init__(self, alpha):
        super().__init__()
        self.alpha =  alpha
    
    def forward(self, x):
        x = torch.mul(x, self.alpha)
        return x

Then use it as:

>>> model = nn.Sequential()
>>> model.add_module(name='activation_output', module=Multiply(10))

>>> model(torch.ones(1))
tensor([10.])