Hi Zimo!



 ZimoNitrome:

Do you happen to know if torch.where breaks the gradent tree?


torch.where() does not break the computation graph – that’s a good
thing about it.
But, furthermore, in this specific case, each individual element of values
only participates in a single the branch of the torch.where() function.
I’m just using it as a more compact way to single out the largest element
in values, rather than indexing with argmax().
Best.
K. Frank