fix autograd bug

brainpy · Sep 18, 2024 · a45583b · a45583b
1 parent 0a1fa85
commit a45583b
Showing 1 changed file with 2 additions and 3 deletions.
diff --git a/brainpy/_src/math/object_transform/autograd.py b/brainpy/_src/math/object_transform/autograd.py
@@ -207,9 +207,8 @@ def __call__(self, *args, **kwargs):
                           *args,
                           **kwargs)
         cache_stack(self.target, stack)
-
-    self._dyn_vars = stack
-    self._dyn_vars.remove_by_id(*[id(v) for v in self._grad_vars])
+      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():