@staticmethod
def backward(ctx, grad_output):
    (k, p, dim) = (*ctx.saved_tensors, ctx.dim)
    grad = grad_output.masked_fill(p.eq(0), 0)
    grad = torch.where(p.ne(0), (grad - (grad.sum(dim, True) / k)), grad)
    return (grad, None)
@staticmethod
def backward(ctx, grad_output):
    (k, p, dim) = (*ctx.saved_tensors, ctx.dim)
    grad = grad_output.masked_fill(p.eq(0), 0)
    grad = torch.where(p.ne(0), (grad - (grad.sum(dim, True) / k)), grad)
    return (grad, None)