[BugFix] Fix typo in weight assignment in PRB (pytorch#2241)

ai4co · Jun 20, 2024 · 1d729e8 · 1d729e8
1 parent 9b1ebb2
commit 1d729e8
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/torchrl/data/replay_buffers/samplers.py b/torchrl/data/replay_buffers/samplers.py
@@ -482,7 +482,7 @@ def sample(self, storage: Storage, batch_size: int) -> torch.Tensor:
             index = torch.where(zero_weight, index - 1, index)
             if (index < 0).any():
                 raise RuntimeError("Failed to find a suitable index")
-            zero_weight = torch.as_tensor(self._sum_tree[index])
+            weight = torch.as_tensor(self._sum_tree[index])
             zero_weight = weight == 0
 
         # Importance sampling weight formula: