Skip to content

Commit b177f22

Browse files
[BugFix] Remove deprecated ray.internal.free calls in RayCollector (#4519)
Co-authored-by: Ben Sprenger <bensprenger1999@gmail.com>
1 parent 3a13e5f commit b177f22

1 file changed

Lines changed: 5 additions & 6 deletions

File tree

  • torchrl/collectors/distributed

‎torchrl/collectors/distributed/ray.py‎

Lines changed: 5 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1317,10 +1317,11 @@ def _sync_iterator(self) -> Iterator[TensorDictBase]:
13171317
out_td = []
13181318
for r in pending_tasks:
13191319
rollouts = ray.get(r)
1320-
ray.internal.free(
1321-
r
1322-
) # should not be necessary, deleted automatically when ref count is down to 0
13231320
out_td.append(rollouts)
1321+
del r
1322+
pending_tasks.clear()
1323+
samples_ready.clear()
1324+
samples_not_ready.clear()
13241325

13251326
# Handle case where replay_buffer is used and rollouts are None
13261327
if out_td[0] is None:
@@ -1439,9 +1440,7 @@ def _async_iterator(self) -> Iterator[TensorDictBase]:
14391440

14401441
# Retrieve single rollouts
14411442
out_td = ray.get(future)
1442-
ray.internal.free(
1443-
[future]
1444-
) # should not be necessary, deleted automatically when ref count is down to 0
1443+
del future, wait_results
14451444

14461445
# Track collected frames - use frames_per_batch since out_td might be None
14471446
# when using RayReplayBuffer (sub-collectors write directly to buffer)

0 commit comments

Comments
 (0)