Update base.py

strongio · Jan 16, 2025 · 1968699 · 1968699
1 parent 3f0e0ed
commit 1968699
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/torchcast/state_space/base.py b/torchcast/state_space/base.py
@@ -461,6 +461,7 @@ def _script_forward(self,
             num_groups = meanu.shape[0]
             if stop_at_last_measured:
                 warn("Ignoring `stop_at_last_measured` since `input` is None.")
+            last_measured_per_group = torch.full((num_groups,), out_timesteps, dtype=torch.int, device=meanu.device)
         else:
             if len(input.shape) != 3:
                 raise ValueError(f"Expected len(input.shape) == 3 (group,time,measure)")