Issue with EAEFNet_Detection EAEF module

[MrNeoBlue]

in your paper, you get the element wise multiplication of tensors after mlp_pooling and create cross_gate and add_gate tensor. In your code, you are not using the diagonal of the result matrix because the argument of torch.diagonal is default——offset=0,dim1=0,dim2=1!

The shape of logits_per is 1x192x192, so torch.diagonal(torch.sigmoid(logits_per)) just extract [:,0,:] of the tensor!

   logits_per = c * rgb_y @ t_y 
   cross_gate = torch.diagonal(torch.sigmoid(logits_per)).reshape(b, c, 1, 1)
   add_gate = torch.ones(cross_gate.shape).cuda() - cross_gate

Anyway, I replaced the torch.diagonal with torch.arange way to get the correct diagonal in order to export the onnx model, the result is nearly the same, with small and medium mAP drop less than 1% and large mAP increases a bit.

If I was wrong, please show me. Many thx!