Not really an nice solution, but if you fill the longtensor data that is shorter than the max index length with repeats, it will work because it will make them the same size. Example from yours modified:
idx = torch.LongTensor( [[0,0,0], [2,2,3], [0,1,4], [0,0,5]])
It will work, but is pretty hacky. I’m sure you could write something to fill the short indices with repeats until it reaches the same length as the longest index array. Maybe someone who knows better can explain if there’s a more elegant workaround.