def forward(self, x: torch.Tensor, mask: torch.Tensor = None) -> torch.Tensor:
        """
        Forward method of the LAMA.

        Args:
            x: tensor of size
                (batch_size, history_len, feature_size)
            mask: mask to use for attention compute

        Returns:
            torch.Tensor: LAMA pooling result
        """
        batch_size, history_len, feature_size = x.shape

        features_list = []
        for pooling_fn in self.groups.values():
            features = pooling_fn(x, mask)
            features_list.append(features)
        x = torch.cat(features_list, dim=1)
        x = x.view(batch_size, -1)

        return x

        Forward method of the LAMA.

        Args:
            x: tensor of size
                (batch_size, history_len, feature_size)
            mask: mask to use for attention compute

        Returns:
            torch.Tensor: LAMA pooling result
        
catalyst/contrib/layers/lama.py