Thanks for the suggestion ! A friend one mine coded a multilinear class as well to deal with that kind of problem
import torch

class MultiLinear(torch.nn.Module):
    def __init__(self, input_size, output_size, nb_estimators, activation=torch.relu):
        super().__init__()
        self.weights = torch.nn.Parameter(torch.randn(nb_estimators, input_size, output_size))
        self.weights.requires_grad = True
        self.activation = activation    def forward(self, x_batch):
        if x_batch.ndim == 2:
            # x_batch: (batch_size, input_size)
            z = torch.einsum("ij,mjl->ilm", x_batch, self.weights)
        elif x_batch.ndim == 3:
            # x_batch: (batch_size, input_size, nb_estimators)
            z = torch.einsum("ijm,mjl->ilm", x_batch, self.weights)
        else:
            raise ValueError("input dims must be 2 or 3")
        output = self.activation(z)
        # output: (batch_size, output_size, nb_estimators)
        return output