Update olmo_model.py

mlc-ai · Nov 23, 2024 · 344404d · 344404d
1 parent ee1a6ee
commit 344404d
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/python/mlc_llm/model/olmo/olmo_model.py b/python/mlc_llm/model/olmo/olmo_model.py
@@ -103,7 +103,9 @@ def __init__(self, config: OlmoConfig):
             self.num_heads * self.head_dim, config.hidden_size, bias=config.attention_bias
         )
 
-    def forward(self, hidden_states: Tensor, paged_kv_cache: PagedKVCache, layer_id: int):# pylint: disable=W0511
+    def forward(
+        self, hidden_states: Tensor, paged_kv_cache: PagedKVCache, layer_id: int
+    ):# pylint: disable=W0511
         d, h_q, h_kv = self.head_dim, self.num_heads, self.num_key_value_heads
         b, s, _ = hidden_states.shape
         qkv = self.qkv_proj(hidden_states)