fix fp16 arg

bagustris · May 30, 2024 · 75c1a2f · 75c1a2f
1 parent c8ace63
commit 75c1a2f
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/nkululeko/models/model_tuned.py b/nkululeko/models/model_tuned.py
@@ -326,7 +326,7 @@ def compute_loss(
             gradient_accumulation_steps=self.accumulation_steps,
             evaluation_strategy="steps",
             num_train_epochs=self.epoch_num,
-            fp16=self.device == "cuda",
+            fp16=self.device != "cpu",
             save_steps=num_steps,
             eval_steps=num_steps,
             logging_steps=num_steps,