Gated MLP Block

class ssm.model.block.GatedMLPBlock(model_dim, hid_dim, activation, beta=1.0)

Bases: Module

Implementation of the Gated MLP Block.

This block combines two linear layers, with one of them passed through a non-linear activation function.

forward(x)

Forward pass of the Gated MLP block.

Parameters:

x (torch.Tensor) – The input tensor.

Returns:

The output tensor.

Return type:

torch.Tensor