[BugFix] Avoid KeyError in slice sampler (for compile)

ghstack-source-id: 6e2a3036f0e50d365387cced50a761b97a47317d Pull Request resolved: #2670
pytorch · Dec 20, 2024 · 21eeca4 · 21eeca4
1 parent f4709c1
commit 21eeca4
Showing 1 changed file with 6 additions and 6 deletions.
diff --git a/torchrl/data/replay_buffers/samplers.py b/torchrl/data/replay_buffers/samplers.py
@@ -1485,13 +1485,13 @@ def _get_index(
                 truncated[seq_length.cumsum(0) - 1] = 1
             index = index.to(torch.long).unbind(-1)
             st_index = storage[index]
-            try:
-                done = st_index[done_key] | truncated
-            except KeyError:
+            done = st_index.get(done_key, default=None)
+            if done is None:
                 done = truncated.clone()
-            try:
-                terminated = st_index[terminated_key]
-            except KeyError:
+            else:
+                done = done | truncated
+            terminated = st_index.get(terminated_key, default=None)
+            if terminated is None:
                 terminated = torch.zeros_like(truncated)
             return index, {
                 truncated_key: truncated,