model samplar_select

allenai · Feb 28, 2024 · cd192bd · cd192bd
1 parent 0129f37
commit cd192bd
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/allenact/algorithms/onpolicy_sync/engine.py b/allenact/algorithms/onpolicy_sync/engine.py
@@ -722,6 +722,9 @@ def collect_step_across_all_task_samplers(
 
         npaused, keep, batch = self.remove_paused(observations)
 
+        if hasattr(self.actor_critic, "sampler_select"):
+            self.actor_critic.sampler_select(keep)
+
         # TODO self.probe(...) can be useful for debugging (we might want to control it from main?)
         # self.probe(dones, npaused)