Minor changes 1

Signed-off-by: Dipankar Sarkar <[email protected]>

Author: qcdipankar

QEfficient/transformers/models/modeling_auto.py

@@ -1415,7 +1415,7 @@ def __init__(
         self.num_layers = model.config.num_hidden_layers
         self.continuous_batching = continuous_batching
         self.model.qaic_config = qaic_config
-        self.pretrained_model_name_or_path = kwargs.pop("pretrained_model_name_or_path", None)
+        self.pretrained_model_name_or_path = kwargs.get("pretrained_model_name_or_path", None)
         self.model, transformed = SpDTransform.apply(self.model, qaic_config, **kwargs)
         self.is_tlm = transformed
 

QEfficient/transformers/models/pytorch_transforms.py

@@ -503,7 +503,7 @@ class SpDTransform:
     @classmethod
     def apply(cls, model: nn.Module, qaic_config: Optional[dict] = None, **kwargs) -> Tuple[nn.Module, bool]:
         transformed = False
-
+        pretrained_model_name_or_path_temp = kwargs.pop("pretrained_model_name_or_path", None)
         if qaic_config is None or (speculative_model_type := qaic_config.get("speculative_model_type")) is None:
             return model, transformed
         elif speculative_model_type not in (
@@ -525,6 +525,7 @@ def apply(cls, model: nn.Module, qaic_config: Optional[dict] = None, **kwargs) -
             raise NotImplementedError(
                 f"model class {model_class} does not yet support returning multiple logits to keep."
             )
+        kwargs["pretrained_model_name_or_path"] = pretrained_model_name_or_path_temp
         return model, transformed