def forward(ctx, a, upper=True): ctx.upper = upper fact = torch.potrf(a, upper) ctx.save_for_backward(fact) return fact