[BugFix] Fix get-related errors (#2361)

pytorch · Aug 5, 2024 · e76d8cb · e76d8cb
1 parent 35a1c5b
commit e76d8cb
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 6 deletions.
diff --git a/torchrl/data/tensor_specs.py b/torchrl/data/tensor_specs.py
@@ -4121,7 +4121,7 @@ def is_in(self, val: Union[dict, TensorDictBase]) -> bool:
         for key, item in self._specs.items():
             if item is None or (isinstance(item, CompositeSpec) and item.is_empty()):
                 continue
-            val_item = val.get(key)
+            val_item = val[key]
             if not item.is_in(val_item):
                 return False
         return True

diff --git a/torchrl/envs/transforms/transforms.py b/torchrl/envs/transforms/transforms.py
@@ -5601,7 +5601,7 @@ def forward(self, tensordict: TensorDictBase) -> TensorDictBase:
             )
         time_dim = time_dim[0] - 1
         for in_key, out_key in zip(self.in_keys, self.out_keys):
-            reward = tensordict.get(in_key)
+            reward = tensordict[in_key]
             cumsum = reward.cumsum(time_dim)
             tensordict.set(out_key, cumsum)
         return tensordict

diff --git a/torchrl/modules/tensordict_module/rnn.py b/torchrl/modules/tensordict_module/rnn.py
@@ -665,7 +665,7 @@ def forward(self, tensordict: TensorDictBase):
         else:
             tensordict_shaped = tensordict.reshape(-1).unsqueeze(-1)
 
-        is_init = tensordict_shaped.get("is_init").squeeze(-1)
+        is_init = tensordict_shaped["is_init"].squeeze(-1)
         splits = None
         if self.recurrent_mode and is_init[..., 1:].any():
             # if we have consecutive trajectories, things get a little more complicated
@@ -679,7 +679,7 @@ def forward(self, tensordict: TensorDictBase):
             tensordict_shaped = _split_and_pad_sequence(
                 tensordict_shaped.select(*self.in_keys, strict=False), splits
             )
-            is_init = tensordict_shaped.get("is_init").squeeze(-1)
+            is_init = tensordict_shaped["is_init"].squeeze(-1)
 
         value, hidden0, hidden1 = (
             tensordict_shaped.get(key, default)
@@ -1410,7 +1410,7 @@ def forward(self, tensordict: TensorDictBase):
         else:
             tensordict_shaped = tensordict.reshape(-1).unsqueeze(-1)
 
-        is_init = tensordict_shaped.get("is_init").squeeze(-1)
+        is_init = tensordict_shaped["is_init"].squeeze(-1)
         splits = None
         if self.recurrent_mode and is_init[..., 1:].any():
             # if we have consecutive trajectories, things get a little more complicated
@@ -1424,7 +1424,7 @@ def forward(self, tensordict: TensorDictBase):
             tensordict_shaped = _split_and_pad_sequence(
                 tensordict_shaped.select(*self.in_keys, strict=False), splits
             )
-            is_init = tensordict_shaped.get("is_init").squeeze(-1)
+            is_init = tensordict_shaped["is_init"].squeeze(-1)
 
         value, hidden = (
             tensordict_shaped.get(key, default)