From adc45266a525e0394180d34a2c07d55580b038dd Mon Sep 17 00:00:00 2001 From: cto Date: Thu, 1 Aug 2024 20:23:36 +0100 Subject: [PATCH] updated wandb config --- simplifine_alpha/train_engine.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/simplifine_alpha/train_engine.py b/simplifine_alpha/train_engine.py index 3cd654c..f49d3c4 100644 --- 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')