Simple example.
You compute 2 samples and u use torch.max between them. Neural network weights will be upgrades with the features corresponding to the max value.  Your example is the same. Backpropagation will flow through those y_hat you are picking, meanwhile those y_hay which you discard will never backpropagate. Therefore your network learns only using y_hat you choose. Despite what you can think trying to constrain the output may not work.
You are slicing the output, choosing some values.  Data flows through them  thus backpropagation too.
About the computational graph, always you use torch functions to generate operate it’s ok.  If you cannot define a tensor based on a previous tensor and torch functions you will probably break the computational graph.  If you cast from one data type to another you also break it.  If you manually assign values of s matrix you break it.
Soz, writing from the phone.