revert

thu-ml · Mar 17, 2022 · 2fd889f · 2fd889f
1 parent 6df9365
commit 2fd889f
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/tianshou/trainer/base.py b/tianshou/trainer/base.py
@@ -240,7 +240,7 @@ def __next__(self) -> Union[None, Tuple[int, Dict[str, Any], Dict[str, Any]]]:
 
             # iterator exhaustion check
             if self.epoch >= self.max_epoch:
-                if self.save_fn:
+                if self.test_collector is not None and self.save_fn:
                     self.save_fn(self.policy)
                 raise StopIteration