Solved by albanD in post #4 
              
              
                Hi, 
You can either design a smoothed version of your loss function that will have non zero gradients. 
Or you can use this piecewise constant function but design a backward to compute something that is not the true gradient but will point in the right direction using a custom autograd.Function.