diff --git a/lzero/model/sampled_efficientzero_model.py b/lzero/model/sampled_efficientzero_model.py
index 58a8197dc..4762e1a71 100644
--- a/lzero/model/sampled_efficientzero_model.py
+++ b/lzero/model/sampled_efficientzero_model.py
@@ -182,6 +182,7 @@ def __init__(
             self.reward_support_size,
             flatten_output_size_for_reward_head,
             downsample,
+            lstm_hidden_size=self.lstm_hidden_size,
             last_linear_layer_init_zero=self.last_linear_layer_init_zero,
             activation=activation,
             norm_type=norm_type