Change in implementation

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

Author: qcdipankar

QEfficient/transformers/models/modeling_auto.py

@@ -1406,10 +1406,9 @@ 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.model, transformed = SpDTransform.apply(self.model, qaic_config, **kwargs)
         self.is_tlm = transformed
-        self.pretrained_model_name_or_path = kwargs.get("pretrained_model_name_or_path", None)
 
         # ---Sampling---
         # Note: SamplerTransform should be applied after all other transforms

QEfficient/transformers/models/pytorch_transforms.py

@@ -500,7 +500,6 @@ class SpDTransform:
     @classmethod
     def apply(cls, model: nn.Module, qaic_config: Optional[dict] = None, **kwargs) -> Tuple[nn.Module, bool]:
         transformed = False
-        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