Added logger for new transform

Signed-off-by: Amit Raj <[email protected]>

Author: quic-amitraj

QEfficient/base/pytorch_transforms.py

@@ -8,6 +8,7 @@
 from typing import Callable, Dict, Tuple, Type
 
 from torch import nn
+from QEfficient.utils.logging_utils import logger
 
 
 class PytorchTransform:
@@ -157,7 +158,7 @@ def apply(cls, model: nn.Module) -> Tuple[nn.Module, bool]:
             if delete_fused_key:
                 del sd[fused_key]
 
-            print(f"[layer {layer_idx:02d}] loaded gate_proj & up_proj from fused tensor  (shape {fused.shape})")
+            logger.info(f"[layer {layer_idx:02d}] loaded gate_proj & up_proj from fused tensor  (shape {fused.shape})")
             transformed = True
         return model, transformed