feat: Don't store data loader in the state

lRomul · Apr 18, 2024 · 8a8903c · 8a8903c
1 parent 4100840
commit 8a8903c
Showing 1 changed file with 2 additions and 3 deletions.
diff --git a/argus/engine/engine.py b/argus/engine/engine.py
@@ -115,7 +115,6 @@ def __init__(self,
             phase_states[self.phase] = self
         self.phase_states = phase_states
         self.logger: logging.Logger = self.model.logger
-        self.data_loader: Optional[Iterable] = None
         self.exception: Optional[BaseException] = None
         self.engine: Optional[Engine] = engine
 
@@ -224,8 +223,7 @@ def run(self, data_loader: Iterable,
             State: An engine state.
 
         """
-        self.state.update(data_loader=data_loader,
-                          epoch=start_epoch,
+        self.state.update(epoch=start_epoch,
                           iteration=0,
                           stopped=False)
 
@@ -250,6 +248,7 @@ def run(self, data_loader: Iterable,
 
                 self.raise_event(Events.EPOCH_COMPLETE)
                 self.state.epoch += 1
+            del data_loader
 
             self.raise_event(Events.COMPLETE)