Hi.

idxs = torch.argmax(prob, dim=-1)
scores = torch.gather(prob, -1, idxs.unsqueeze(-1))

The model exported out onnx using GatherElements,
is there a way to make it export in Gather rather than GatherElements?