


out = fn(y_true, y_preds)
> output is ((grads w.r.t y_true, grads w.r.t y_preds), (y_pred, loss_per_sample))



Note:
 Using PyTorch "torch.no_grad" together with "grad".Case 1: Using
 "torch.no_grad" inside a function:

 >>> def f(x):
 >>>     with torch.no_grad():
 >>>         c = x ** 2
 >>>     return x - c

 In this case, "grad(f)(x)" will respect the inner
 "torch.no_grad".Case 2: Using "grad" inside "torch.no_grad"
 context manager:

 >>> with torch.no_grad():
 >>>     grad(f)(x)

 In this case, "grad" will respect the inner "torch.no_grad", but
 not the outer one. This is because "grad" is a "function
 transform": its result should not depend on the result of a
 context manager outside of "f".
