fix a minor bug in prednet

wp8733684 · Jun 3, 2023 · 389a7d3 · 389a7d3
1 parent 1c48ff0
commit 389a7d3
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/openstl/models/prednet_model.py b/openstl/models/prednet_model.py
@@ -173,7 +173,8 @@ def forward(self, A0_withTimeStep, initial_states=None, extrapolation=False):
 
         '''
         if initial_states is None:
-            initial_states = get_initial_states((1, 10, 1, 64, 64),
+            T, C, H, W = self.args.in_shape
+            initial_states = get_initial_states((1, T, C, H, W)),
                                                 self.row_axis, self.col_axis, self.num_layers, self.R_stack_sizes, self.stack_sizes, self.channel_axis, self.args.device)
         A0_withTimeStep = A0_withTimeStep.transpose(0, 1)
         num_timesteps = A0_withTimeStep.shape[0]