Everything works fine by adding         :
self.clf = nn.Sequential(
nn.Conv1d(in_channels=M, out_channels=32, kernel_size=1, stride = 1),
nn.LeakyReLU(0.2),
nn.Conv1d(in_channels=32, out_channels=16, kernel_size=1, stride=1),
nn.LeakyReLU(0.2),
nn.Conv1d(in_channels=16, out_channels=2, kernel_size=1, stride=1),
)
The only thing that i had to modify was  at the forward function :
    z = z.unsqueeze(-1)
    logit = self.clf(z)
    logit = logit.squeeze(-1)

Where z was the output of an embedding + attention pooling stage
Thank’s for the help