diff --git a/flagai/auto_model/auto_loader.py b/flagai/auto_model/auto_loader.py
index 77831231..51d14c60 100755
--- a/flagai/auto_model/auto_loader.py
+++ b/flagai/auto_model/auto_loader.py
@@ -305,9 +305,11 @@ def __init__(self,
                         math.ceil(model_max_length / orig_ctx_len))
                     config.rope_scaling = {"type": "linear", "factor": scaling_factor}
                 config.use_cache = False
-                model = AquilaForCausalLM.from_pretrained(download_path,
+                model = AquilaForCausalLM.from_pretrained(download_path,low_cpu_mem_usage=low_cpu_mem_usage,
                                                         **kwargs)
-
+                # from accelerate import load_checkpoint_and_dispatch
+                # model = load_checkpoint_and_dispatch(
+                #     model, download_path, device_map="auto", no_split_module_classes=["AquilaDecoderLayer"])
             tokenizer = AutoTokenizer.from_pretrained(download_path)
             self.model = model 
             self.tokenizer = tokenizer