Hi,
The problem is that the output of count_nonzero is a Tensor of type int64. So it cannot require gradients (only continuous types can).
So you want to add a output = output.float() after the count_nonzero to make sure the output of your custom Function can require gradients.