fix autograd bug

brainpy · Sep 18, 2024 · 10043a0 · 10043a0
1 parent a45583b
commit 10043a0
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/brainpy/_src/math/object_transform/autograd.py b/brainpy/_src/math/object_transform/autograd.py
@@ -210,9 +210,9 @@ def __call__(self, *args, **kwargs):
       self._dyn_vars = stack
       self._dyn_vars.remove_by_id(*[id(v) for v in self._grad_vars])
 
-    # if not the outermost transformation
-    if not stack.is_first_stack():
-      return self._return(rets)
+      # if not the outermost transformation
+      if not stack.is_first_stack():
+        return self._return(rets)
 
     rets = self._transform(
       [v.value for v in self._grad_vars],  # variables for gradients