Hi, I see the topic and I have the same problem.
My code is as follow:
Each of my inputs are a tensor of size (512, 20).
class MyModel(nn.Module):
  def __init__(self, input_dim: int = 512, hidden_dim: int = 50, num_layers: int = 4, dropout: float = 0.5):
    super(MyModel, self).__init__()
    self.input_dim = input_dim
    self.hidden_dim = hidden_dim
    self.num_layers = num_layers
    self.dropout = dropout
    self.fc = nn.Linear(in_features=self.num_layers * self.hidden_dim, out_features=1)
    self.bn = nn.BatchNorm1d(num_features=self.num_layers * self.hidden_dim)

    convs_list = [nn.Conv1d(in_channels=self.input_dim, out_channels=self.hidden_dim, kernel_size=i+1) for i in range(self.num_layers)]
    self.convs = nn.ModuleList(modules=convs_list)

  def reset_parameters(self):
        for conv in self.convs:
            conv.reset_parameters()
    
  def forward(self, x):
    out = []
    for i in range(self.num_layers):
      x_i = self.convs[i](x)
      x_i = F.leaky_relu(x_i, negative_slope=0.02)
      x_i = F.dropout(x_i, self.dropout, self.training)
      x_i = x_i.mean(dim=0)
      out.append(x_i)
    out = torch.cat(out, dim=1)
    out = self.bn(out)
    out = self.fc(out)
    return out