diff --git a/transformer_lens/pretrained/weight_conversions/olmoe.py b/transformer_lens/pretrained/weight_conversions/olmoe.py index f673429cc..8a9a2967b 100644 --- a/transformer_lens/pretrained/weight_conversions/olmoe.py +++ b/transformer_lens/pretrained/weight_conversions/olmoe.py @@ -3,6 +3,7 @@ from transformer_lens.HookedTransformerConfig import HookedTransformerConfig + def convert_olmoe_weights(olmoe, cfg: HookedTransformerConfig): state_dict = {}