Dear Lightning-Community,
the following callback does not create an autograd graph, as x2.grad_fn is None:
class AdditionalEvaluations(L.Callback):
def on_test_end(self, trainer, module):
with torch.enable_grad():
x = Tensor([1.]).requires_grad_(True)
print(“\n”)
print("tensor: ", x, "requires_grad: ", x.requires_grad, “grad_fn”, x.grad_fn)
x2 = x**2+1.
print("tensor: ", x2, "requires_grad: ", x2.requires_grad, “grad_fn”, x2.grad_fn)
print("is_grad_enabled: ", torch.is_grad_enabled())
The following notebook implements the BoringModel showing this behaviour: Google Colab
Is there any way to build a graph within the on_test_end callback?
Thank you!
msurner