updated wandb config

simplifine-llm · Aug 1, 2024 · adc4526 · adc4526
1 parent 430fd59
commit adc4526
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/simplifine_alpha/train_engine.py b/simplifine_alpha/train_engine.py
@@ -860,6 +860,7 @@ def formatting_prompts_func(example):
             )
 
     # initialize the model
+    device, device_name = init_device()
     if device_name == 'cuda':
         model = AutoModelForCausalLM.from_pretrained(model_name, 
                                                  token = hf_token,
@@ -875,7 +876,7 @@ def formatting_prompts_func(example):
     if use_peft:
         model = get_peft_model(model, peft_config=peft_config)
 
-    device, device_name = init_device()
+
     if torch.cuda.device_count() > 1:
         if ddp and zero:
             raise ValueError('Zero optimization and DDP cannot be used together')