Pytorch assigning to results of gather

[sachinruk (Sachinthaka Abeywardana)]

Hi all, just cross posting from SO. Wondering if there was a more efficient solution: pytorch assigning to results of gather - Stack Overflow