should be fine?

Parry-Parry · Nov 6, 2024 · 4f377d5 · 4f377d5
1 parent 8a3211e
commit 4f377d5
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/rankers/modelling/cat.py b/rankers/modelling/cat.py
@@ -192,8 +192,8 @@ def to_pyterrier(self) -> "pt.Transformer":
     @classmethod
     def from_pretrained(cls, model_dir_or_name : str, num_labels=2, config=None, **kwargs) -> "Cat":
         """Load model from a directory"""
-        config = cls.config_class.from_pretrained(model_dir_or_name) if config is None else config
-        model = cls.architecture_class.from_pretrained(model_dir_or_name, num_labels=num_labels, config=config, **kwargs)
+        config = cls.config_class.from_pretrained(model_dir_or_name, num_labels=num_labels) if config is None else config
+        model = cls.architecture_class.from_pretrained(model_dir_or_name, config=config, **kwargs)
         tokenizer = AutoTokenizer.from_pretrained(model_dir_or_name)
         return cls(model, tokenizer, config)