Use column variable in train

louisbrulenaudet · Apr 3, 2024 · 34f479e · 34f479e
1 parent 9663a2f
commit 34f479e
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/src/tsdae/tsdae.py b/src/tsdae/tsdae.py
@@ -856,7 +856,7 @@ def train(
             # Constrain dataset
             constrained_data = self.constrain_dataset(
                 shuffled_dataset=shuffled_dataset,
-                column="output",
+                column=column,
                 splitter=splitter,
                 num_to_keep=num_to_keep,
                 total_sentences=total_sentences,
@@ -904,4 +904,4 @@ def train(
 
         except Exception as e:
             raise Exception(f"Error in train: {str(e)}")
-            return None
+            return None