@staticmethod
def backward(ctx, grad_output):
    (in_shape,) = ctx.saved_tensors
    return grad_output.reshape(in_shape)
@staticmethod
def backward(ctx, grad_output):
    (in_shape,) = ctx.saved_tensors
    return grad_output.reshape(in_shape)