grumble

stanford-crfm · Nov 16, 2024 · 122365e · 122365e
1 parent f421444
commit 122365e
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/levanter/models/whisper.py b/src/levanter/models/whisper.py
@@ -447,7 +447,7 @@ def _state_dict_key_map(self) -> Dict[str, Optional[str]]:
         return {"transformer": None, "embeddings": None}
 
 
-class WhisperModel(ModelWithHfSerializationMixin[WhisperConfig]):
+class WhisperModel(eqx.Module, ModelWithHfSerializationMixin[WhisperConfig]):
     encoder: WhisperEncoder
     decoder: WhisperDecoder