def METHOD_NAME(self, x: torch.Tensor) -> torch.Tensor:
    """
    x: A tensor of dimension (B, L, M) where
    B is the batch size, L is the length of the sequence
    """
    x = self.initial_layer(x)
    x = x.permute(0, 2, 1)
    x = self.output_layer(x)
    if self.probabilistic:
        mean = x[..., 0][..., None]
        std = torch.clamp(x[..., 1][..., None], min=0.01)
        return torch.distributions.Normal(mean, std)
    else:
        return x.view(-1, self.output_len)