replace init_empty_weights with torch.device(meta)

lchu-ibm · Aug 1, 2023 · c8d4f38 · c8d4f38
1 parent d8a81bb
commit c8d4f38
Showing 1 changed file with 1 addition and 3 deletions.
diff --git a/llama_finetuning.py b/llama_finetuning.py
@@ -42,8 +42,6 @@
     get_policies  
 )
 
-from accelerate import init_empty_weights
-
 from utils.dataset_utils import get_preprocessed_dataset
 
 from utils.config_utils import (
@@ -107,7 +105,7 @@ def main(**kwargs):
             )
         else:
             llama_config = LlamaConfig.from_pretrained(train_config.model_name)
-            with init_empty_weights():
+            with torch.device("meta"):
                 model = LlamaForCausalLM(llama_config)
     else:
         model = LlamaForCausalLM.from_pretrained(