The solution works beautifully but if you know a way to achieve this without creating a new tensor (arange), it would be even better.
The reason I needed it is because I am trying to use pytorch/glow project but the do not support dynamically creating new tensors (see glow/issues/3932 5)