Skip to content

Commit bd3d4f3

Browse files
committed
[Feature] Macro-actions for LLMs
ghstack-source-id: d784e33 Pull Request resolved: #2831
1 parent b204dc8 commit bd3d4f3

File tree

1 file changed

+4
-1
lines changed

1 file changed

+4
-1
lines changed

torchrl/modules/tensordict_module/actors.py

+4-1
Original file line numberDiff line numberDiff line change
@@ -2371,7 +2371,10 @@ def forward(
23712371
action_entry = parent_td.get(action_key_orig[-1], None)
23722372
if action_entry is None:
23732373
raise self._NO_INIT_ERR
2374-
if self.n_steps is not None and action_entry.shape[parent_td.ndim] != self.n_steps:
2374+
if (
2375+
self.n_steps is not None
2376+
and action_entry.shape[parent_td.ndim] != self.n_steps
2377+
):
23752378
raise RuntimeError(
23762379
f"The action's time dimension (dim={parent_td.ndim}) doesn't match the n_steps argument ({self.n_steps}). "
23772380
f"The action shape was {action_entry.shape}."

0 commit comments

Comments
 (0)