Update modeling_hymba.py

modeling_hymba.py

@@ -1714,7 +1714,7 @@ class HymbaBlock(nn.Module):
     
                     if ssm_state is not None and cache_params is not None:
                         cache_params.ssm_states[self.layer_idx].copy_(ssm_state)
-            if use_precomputed_states:
+            if use_precomputed_states and self.layer_idx==31:
             # except Exception as e:
                 print("\n\n\n\n")
                 # print(e)