wip

datamol-io · Jul 29, 2024 · 04b96b5 · 04b96b5
1 parent 3a2e40d
commit 04b96b5
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/safe/trainer/cli.py b/safe/trainer/cli.py
@@ -331,7 +331,8 @@ def compute_metrics(eval_preds):
 
     if model_args.include_descriptors:
         training_args.label_names = ["labels", "mc_labels"]
-    else:
+
+    if training_args.label_names is None:
         training_args.label_names = ["labels"]
     # update dispatch_batches in accelerator
     training_args.accelerator_config.dispatch_batches = data_args.streaming is not True