Sorry for bringing this up again, but has anybody an idea how to implement the backward path?
or is it simply passing the gradient through, where the mask is one? like this:

    @staticmethod
    def backward(ctx, grad_output):
        mask_active, = ctx.saved_tensors
        grad_input = grad_output.clone()
        grad_input[~mask_active] = 0
        return grad_input