Avoid stream sync in softmax

[spinjo]

Bug in `torch_geometric.utils.softmax

If the output_size argument is not provided in torch.repeat_interleave, then the output_size is calculated on the fly, at the cost of a stream synchronization, see https://docs.pytorch.org/docs/stable/generated/torch.repeat_interleave.html. But if the torch_geometric.utils.softmax function gets the index argument in addition to ptr, we get output_size for free - so we should use it.

Proposed fix

When ptr and index are both provided in torch_geometric.utils.softmax, use index to set the output_size in torch.repeat_interleave. Added a unit test for this case.

[spinjo]

PR was open for 5 months, but should be an easy one to merge. Let me know if I can help with anything.

@akihironitta @rusty1s @wsad1