hello @albanD ,
I wrote following program for the same (no concat in this, only computing scores), is this a valid test, or am I making some mistake?
class A(nn.Module):
  def __init__(self):
    super().__init__()
    self.embed = nn.Embedding(50, 10)
  def forward(self, inp_matrix):
    inp_matrix = inp_matrix.reshape(50)
    self.scores = torch.zeros(50)
    for i in range(len(inp_matrix)):
      for j in range(len(inp_matrix)):
        self.scores[i] += abs(((self.embed(inp_matrix[i].long())@(self.embed(inp_matrix[i].long()) + self.embed(inp_matrix[j].long())).t()))) # I use formula (A)*(A+B)
    self.scores = self.scores/(self.scores.sum().expand_as(self.scores)) # I divide values by fixed number, but this should also be learnt
    # print(self.scores)
    return self.scores

model = A()
loss_func = nn.CrossEntropyLoss() # .to('cuda')
optim = torch.optim.SGD(model.parameters(), lr=5) # I use SGD, for technique to be a better fit for SGD

model.train()
for i in range(1000):
  optim.zero_grad()
  input = torch.randperm(50).float()
  zeros = torch.zeros(50)
  input.requires_grad = True
  target = torch.where(input==15, torch.tensor(1.), zeros)
  print('input', input, 'target', target)
  input = input.reshape(1, 50)
  # print(input.shape)
  # print('input', input, '\n', 'target', target)
  x = model(input)
  x = x.reshape(1, 50)
  y = torch.argmax(x)
  print('x', x, 'target', torch.tensor([torch.argmax(target).item()]))
  loss = loss_func(x, torch.tensor([torch.argmax(target).item()]))
  loss.backward()
  # for p in model.parameters():
    # print(p.grad)
  print(list(model.parameters())[0].sum())
  optim.step()
  print(list(model.parameters())[0].sum())
  output = input.reshape(50)
  for j in range(len(output)):
    # print(j, y, input[j], input[y])
    if output[j] < output[y]:
      output[j] = 0.
  print('output', output)
  print(loss)

I see loss decreasing, and output, is 0 for values below 15.