diff --git a/deepxde/gradients/gradients.py b/deepxde/gradients/gradients.py
index 742796b50..b819405a5 100644
--- a/deepxde/gradients/gradients.py
+++ b/deepxde/gradients/gradients.py
@@ -19,7 +19,7 @@ def jacobian(ys, xs, i=None, j=None):
       computation.
 
     Args:
-        ys: Output Tensor of shape (batch_size, dim_y).
+        ys: Output Tensor of shape (batch_size, dim_y) or (batch_size1, batch_size2, dim_y).
         xs: Input Tensor of shape (batch_size, dim_x).
         i (int or None): `i`th row. If `i` is ``None``, returns the `j`th column
             J[:, `j`].