rollback loader.py

Author: lwj2015

models/loader.py

@@ -5,6 +5,8 @@
 import warnings
 from typing import Optional
 
+import torch
+
 from configs.model_config import DEVICE
 from models.adapter import get_llm_model_adapter
 from models.compression import compress_module
@@ -52,7 +54,6 @@ class ModelLoader(metaclass=Singleton):
     def __init__(self, model_path) -> None:
         self.device = DEVICE
         self.model_path = model_path
-        import torch
         self.kwargs = {
             "torch_dtype": torch.float16,
             "device_map": "auto",
@@ -67,7 +68,6 @@ def loader(
         cpu_offloading=False,
         max_gpu_memory: Optional[str] = None,
     ):
-        import torch
         if self.device == "cpu":
             kwargs = {"torch_dtype": torch.float32}