Solved by albanD in post #2 
              
              
                Hi, 
You can. write a custom Function whose backward will be the identity while the forward is your torch.where. 
You can find here how to write such Function.