issue_211

pyadjoint/block_variable.py

@@ -21,6 +21,8 @@ def __init__(self, output):
         self.creation_timestep = -1
         # The timestep during which this variable was last used as an input.
         self.last_use = -1
+        # Whether this variable is adjoint dependency.
+        self.is_adjoint_dependency = False
 
     def add_adj_output(self, val):
         if self.adj_value is None:
@@ -82,6 +84,7 @@ def will_add_as_output(self):
         if tape._eagerly_checkpoint_outputs:
             self.save_output()
         tape.add_to_adjoint_dependencies(self, self.last_use - 1)
+        self.is_adjoint_dependency = True
 
     def __str__(self):
         return str(self.output._ad_str)

pyadjoint/checkpointing.py

@@ -361,7 +361,7 @@ def _(self, cp_action, progress_bar, functional=None, **kwargs):
 
                 # Handle the case for SingleMemoryStorageSchedule
                 if isinstance(self._schedule, SingleMemoryStorageSchedule):
-                    if step > 1 and var not in self.tape.timesteps[step - 1].adjoint_dependencies:
+                    if not var.is_adjoint_dependency:
                         var._checkpoint = None
                     continue