[BugFix] Fix update_priority generic signature for Samplers (pytorch#…

…2252)
ai4co · Jun 27, 2024 · 353c897 · 353c897
1 parent 670a8cf
commit 353c897
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/torchrl/data/replay_buffers/samplers.py b/torchrl/data/replay_buffers/samplers.py
@@ -57,7 +57,11 @@ def extend(self, index: torch.Tensor) -> None:
         return
 
     def update_priority(
-        self, index: Union[int, torch.Tensor], priority: Union[float, torch.Tensor]
+        self,
+        index: Union[int, torch.Tensor],
+        priority: Union[float, torch.Tensor],
+        *,
+        storage: Storage | None = None,
     ) -> dict | None:
         return