Yes, indeed, it destroys the chain. I guess in many contexts this is not a problem.
I think a context where you want to optimize on the mask is more likely to be some kind of RL problem where you have a discrete action space.
Variable.nonzero is not implemented yet as discussed in the link below. However, if it was, I wonder how the backward would be implemented, since it outputs indices…

  
    
    
    Indexing a Variable with a mask generated from another Variable 
  
  
    Four months passed, we still can not use nonzero() for Variables…
  


See an interesting discussion here too:

  
    
    
    Find non zero elements in a tensor 
  
  
    Yes, of course your total loss L is (piecewise) differentiable. It can be more formally defined as: 
L = sum( Li ) / sum( 1{Li ≠ 0} ), 
where 1{c} is the indicator function (which is 1 when c is true and 0 otherwise). 
Clearly, the function f ( Li ) = 1{Li ≠ 0} has derivative equal to 0 everywhere, except at Li = 0, where the derivative does not exist. In practice, you may assume that it is 0 everywhere. 
Your biggest concern should be ensuring that you have no problem in using the tensor losses…