Note that torch.max and your approach are a bit different.
While torch.max works on a tensor and will pass the gradient to the max value (as you’ve explained), your approach would create different computation graphs.
Since PyTorch creates the computation graph (which is used to calculate the gradients in the backward pass) dynamically, you can use plain Python conditions, loops, etc. in your code.