Fixed issue when fitting per episodes

MushroomRL · Dec 8, 2023 · b61c679 · b61c679
1 parent 03c180c
commit b61c679
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/mushroom_rl/core/_impl/vectorized_core_logic.py b/mushroom_rl/core/_impl/vectorized_core_logic.py
@@ -63,10 +63,13 @@ def after_step(self, last):
         self._current_episodes_counter += completed
         self._episodes_progress_bar.update(completed)
 
-    def after_fit(self):
+    def after_fit_vectorized(self, last):
         super().after_fit()
         if self._n_episodes_per_fit is not None:
             self._running_envs = self._array_backend.zeros(self._n_envs, dtype=bool)
+            return self._array_backend.ones(self._n_envs, dtype=bool)
+        else:
+            return last
 
     def _reset_counters(self):
         super()._reset_counters()