Fix cache in UniversalTransformerCache

modeling_ouro.py

@@ -146,21 +146,22 @@ class UniversalTransformerCache(Cache):
     ) -> tuple[torch.Tensor, torch.Tensor]:
         if layer_idx < 0:
             raise ValueError(f"layer_idx must be non-negative, got {layer_idx}")
-
+        
         if self.max_cache_size is not None and layer_idx >= self.max_cache_size:
-            raise IndexError(
-                f"Cache index {layer_idx} exceeds configured max_cache_size={self.max_cache_size}. "
-                "Check total_ut_steps and num_hidden_layers."
-            )
-
+            self.max_cache_size = layer_idx + 1
+        
+        while len(self.key_cache) <= layer_idx:
+            self.key_cache.append(None)
+            self.value_cache.append(None)
+        
         # Expand cache storage so the requested index is available.
         while len(self.key_cache) <= layer_idx:
             self.key_cache.append(None)
             self.value_cache.append(None)
-
+        
         cached_key = self.key_cache[layer_idx]
         cached_value = self.value_cache[layer_idx]
-
+        
         if cached_key is None:
             self.key_cache[layer_idx] = key_states
             self.value_cache[layer_idx] = value_states