diff --git a/utils/classification.py b/utils/classification.py
index df41ef1..e386e38 100644
--- a/utils/classification.py
+++ b/utils/classification.py
@@ -299,7 +299,7 @@ def train_classifying_direction(
                 **kwargs,
             )
             if test_data is None:
-                test_data = train_data
+                test_line = train_line
             else:
                 test_line, _, _, _ = fitting_method(
                     test_data, test_pos, test_layer,