amend

pytorch · vmoens · May 31, 2024 · Apr 30, 2024 · May 26, 2024 · May 26, 2024
commit c13549f685ad92cd347f68ac888b72ed4e52b90b
diff --git a/test/test_collector.py b/test/test_collector.py
@@ -560,6 +560,7 @@ def env_fn(seed):
         total_frames=20000,
         device="cpu",
     )
+    assert collector._use_buffers
     for i, d in enumerate(collector):
         if i == 0:
             b1 = d
@@ -587,14 +588,18 @@ def env_fn(seed):
             b2c = d
         else:
             break
+    assert ccollector._use_buffers
     assert d.names[-1] == "time"
-    with pytest.raises(AssertionError):
-        assert_allclose_td(b1c, b2c)
 
-    assert_allclose_td(b1c, b1)
-    assert_allclose_td(b2c, b2)
+    try:
+        with pytest.raises(AssertionError):
+            assert_allclose_td(b1c, b2c)
 
-    ccollector.shutdown()
+        assert_allclose_td(b1c, b1)
+        assert_allclose_td(b2c, b2)
+    finally:
+        ccollector.shutdown()
+        del ccollector
 
 
 @pytest.mark.skipif(not _has_gym, reason="gym library is not installed")

diff --git a/torchrl/collectors/collectors.py b/torchrl/collectors/collectors.py
@@ -2424,10 +2424,10 @@ def _get_from_queue(self, timeout=None) -> Tuple[int, int, TensorDictBase]:
                 data, idx = new_data
                 self.out_tensordicts[idx] = data
                 if use_buffers is None and j > 0:
-                    self._use_buffers = False
+                    use_buffers = self._use_buffers = False
             except TypeError:
                 if use_buffers is None:
-                    self._use_buffers = True
+                    use_buffers = self._use_buffers = True
                     idx = new_data
                 else:
                     raise