Fix labels missing indexing

psaegert · Feb 9, 2025 · c76b637 · c76b637
1 parent a2479e4
commit c76b637
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/src/flash_ansr/compat/evaluation_nesymres.py b/src/flash_ansr/compat/evaluation_nesymres.py
@@ -112,7 +112,7 @@ def evaluate(
                 results_dict['n_support'].append([batch['x_tensors'].shape[1] // 2] * batch['x_tensors'].shape[0])
 
                 # Create the labels for the next token prediction task (i.e. shift the batch['input_ids'] by one position to the right)
-                labels = batch['labels'].clone()
+                labels = batch['labels'][0].clone()
                 labels_decoded = expression_space.tokenizer.decode(labels.tolist(), special_tokens='<num>')
 
                 # TODO: For different datasets, sort unused dimensions to the end

diff --git a/src/flash_ansr/compat/evaluation_pysr.py b/src/flash_ansr/compat/evaluation_pysr.py
@@ -132,7 +132,7 @@ def evaluate(
                 results_dict['n_support'].append([x_tensor.shape[1] // 2] * x_tensor.shape[0])
 
                 # Create the labels for the next token prediction task (i.e. shift the input_ids by one position to the right)
-                labels = batch['labels'].clone()
+                labels = batch['labels'][0].clone()
                 labels_decoded = expression_space.tokenizer.decode(batch['labels'].tolist(), special_tokens='<num>')
 
                 # TODO: For different datasets, sort unused dimensions to the end