update code

Author: Jerry-Jzy

deepxde/gradients/jacobian.py

@@ -114,7 +114,9 @@ def __call__(self, ys, xs, i=None, j=None):
         #     x = torch.from_numpy(x)
         #     x.requires_grad_()
         #     f(x)
-        if backend_name in ["tensorflow.compat.v1", "tensorflow"]:
+        if backend_name == "tensorflow.compat.v1":
+            key = (ys.ref(), xs.ref())
+        elif backend_name == "tensorflow":
             if config.autodiff == "reverse":
                 key = (ys.ref(), xs.ref())
             elif config.autodiff == "forward":