From a45583b3046b14434a8c45c36f6468fb9103d4d0 Mon Sep 17 00:00:00 2001 From: Chaoming Wang Date: Wed, 18 Sep 2024 15:34:30 +0800 Subject: [PATCH] fix autograd bug --- brainpy/_src/math/object_transform/autograd.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/brainpy/_src/math/object_transform/autograd.py b/brainpy/_src/math/object_transform/autograd.py index 1a482cc8..f89f33fa 100644 --- 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():