amend

pytorch · vmoens · May 31, 2024 · Apr 30, 2024 · May 26, 2024 · May 26, 2024
commit d1a052fefe0846030d510225fe310a1710b310e3
diff --git a/test/test_collector.py b/test/test_collector.py
@@ -1840,6 +1840,7 @@ def test_set_truncated(collector_cls):
             break
     finally:
         collector.shutdown()
+        del collector
 
 
 class TestNestedEnvsCollector:

diff --git a/torchrl/_utils.py b/torchrl/_utils.py
@@ -118,6 +118,7 @@ def _check_for_faulty_process(processes):
             for _p in processes:
                 if _p.is_alive():
                     _p.terminate()
+                    _p.close()
         if terminate:
             break
     if terminate:

diff --git a/torchrl/collectors/collectors.py b/torchrl/collectors/collectors.py
@@ -809,9 +809,6 @@ def _set_truncated_keys(self):
             self._truncated_keys = [
                 key for key in self.env.done_keys if _ends_with(key, "truncated")
             ]
-            self._truncated_keys = self._truncated_keys + [
-                unravel_key(("next", key)) for key in self._truncated_keys
-            ]
 
     @classmethod
     def _get_devices(
@@ -1779,8 +1776,8 @@ def _shutdown_main(self) -> None:
             self.queue_out.close()
             for pipe in self.pipes:
                 pipe.close()
-                for proc in self.procs:
-                    proc.join(1.0)
+            for proc in self.procs:
+                proc.join(1.0)
         finally:
             import torchrl