element 0 of tensors does not require grad and does not have a grad_fn in "test_step" and "validation_step" about pytorch-lightning HOT 4 OPEN

    with torch.enable_grad():
        N = 5
        f = lambda x: x ** 2
        x = torch.randn(N, requires_grad=True)
        # print(x.requires_grad)
        y = f(x)
        I_N = torch.eye(N)
        print(torch.is_inference_mode_enabled()) # False in train and validation, but True in test.
        # Sequential approach
        jacobian_rows = [torch.autograd.grad(y, x, v, retain_graph=True,allow_unused=True)[0]
                        for v in I_N.unbind()]
        jacobian = torch.stack(jacobian_rows)

    return jacobian

with torch.enable_grad():
    x= torch.randn(3,10).requires_grad_(True)
    y = x+1
    print(f"{torch.is_inference_mode_enabled()=}")
    print(f"{x.requires_grad=}")
    print(f"{y.requires_grad=}")

with torch.inference_mode():
    with torch.enable_grad():
        x= torch.randn(3,1).requires_grad_(True)
        y = x+1
        print(f"{torch.is_inference_mode_enabled()=}")
        print(f"{x.requires_grad=}")
        print(f"{y.requires_grad=}")

with torch.inference_mode():
    x= torch.randn(3,10).requires_grad_(True)
    y = x+1
    print(f"{torch.is_inference_mode_enabled()=}")
    print(f"{x.requires_grad=}")
    print(f"{y.requires_grad=}")
    with torch.inference_mode(mode=False):
        # with torch.enable_grad():
            x= torch.randn(3,10).requires_grad_(True)
            y = x+1
            print(f"{torch.is_inference_mode_enabled()=}")
            print(f"{x.requires_grad=}")
            print(f"{y.requires_grad=}")

torch.is_inference_mode_enabled()=True
x.requires_grad=True
y.requires_grad=False
torch.is_inference_mode_enabled()=False
x.requires_grad=True
y.requires_grad=True