For the original poster’s question of how to use backward within a function, the solution is to wrap the calculation in with torch.enable_grad(): as in
  with torch.enable_grad():
            theta = solve_logistic_regression(X, y, lamb)

(and, of course, return something reasonable).
You can think of the forward of an autograd.Function to be inside a with torch.no_grad(): block, so you need to reenable tracking.
Best regards
Thomas
P.S.: As a side note, when posting your question in multiple venues (which you should only in very rare cases), it is a great service to those answering and following along if you add references to your cross posting. I thought I had seen the exact question before, and no, it’s not a bug. 24