Undo erroneous Hessian modification

tttc3 · tttc3 · commit f4fdc0cd5f03 · 2023-09-24T20:52:12.000+01:00
diff --git a/deepxde/gradients.py b/deepxde/gradients.py
@@ -203,7 +203,7 @@ def __init__(self, y, xs, component=None, grad_y=None):
         if backend_name in ["tensorflow.compat.v1", "tensorflow", "pytorch", "paddle"]:
             dim_y = y.shape[1]
         elif backend_name == "jax":
-            dim_y = y[0].shape[1]
+            dim_y = y[0].shape[0]
 
         if dim_y > 1:
             if component is None:
