[BugFix] requested_frames_per_batch in distributed collectors

Vincent Moens · Vincent Moens · commit 408cf7d04705 · 2024-11-18T17:10:50.000Z
ghstack-source-id: 49289de Pull Request resolved: #2579
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
