[BugFix] requested_frames_per_batch in distributed collectors

ghstack-source-id: 49289de6956460d9aed13d982eb8003eafc35118 Pull Request resolved: #2579
pytorch · Nov 18, 2024 · 408cf7d · 408cf7d
1 parent db7f08d
commit 408cf7d
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 0 deletions.
diff --git a/torchrl/collectors/distributed/generic.py b/torchrl/collectors/distributed/generic.py
@@ -448,6 +448,7 @@ def __init__(
         self.policy_weights = policy_weights
         self.num_workers = len(create_env_fn)
         self.frames_per_batch = frames_per_batch
+        self.requested_frames_per_batch = frames_per_batch
 
         self.device = device
         self.storing_device = storing_device

diff --git a/torchrl/collectors/distributed/rpc.py b/torchrl/collectors/distributed/rpc.py
@@ -304,6 +304,7 @@ def __init__(
         self.policy_weights = policy_weights
         self.num_workers = len(create_env_fn)
         self.frames_per_batch = frames_per_batch
+        self.requested_frames_per_batch = frames_per_batch
 
         self.device = device
         self.storing_device = storing_device

diff --git a/torchrl/collectors/distributed/sync.py b/torchrl/collectors/distributed/sync.py
@@ -315,6 +315,7 @@ def __init__(
         self.policy_weights = policy_weights
         self.num_workers = len(create_env_fn)
         self.frames_per_batch = frames_per_batch
+        self.requested_frames_per_batch = frames_per_batch
 
         self.device = device
         self.storing_device = storing_device