use .shape instead of len() when do shape eval

microsoft · Jan 8, 2025 · 6d8f476 · WinstonLiyt · Jan 9, 2025 · 6d8f476
1 parent 56d57ac
commit 6d8f476
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/rdagent/components/coder/data_science/feature/eval_tests/feature_test.py b/rdagent/components/coder/data_science/feature/eval_tests/feature_test.py
@@ -19,7 +19,7 @@
 
 
 # Validate the conditions mentioned in the docstring
-assert len(X_test) == len(test_ids), "Mismatch in length of test images and test IDs"
-assert len(X) == len(y), "Mismatch in length of training images and labels"
+assert X_test.shape[0] == test_ids.shape[0], "Mismatch in length of test images and test IDs"
+assert X.shape[0] == y.shape[0], "Mismatch in length of training images and labels"
 
 print("Feature Engineering test passed successfully. Length of test images matches length of test IDs.")
diff --git a/rdagent/components/coder/data_science/raw_data_loader/eval_tests/data_loader_test.py b/rdagent/components/coder/data_science/raw_data_loader/eval_tests/data_loader_test.py
@@ -16,8 +16,8 @@
 X, y, X_test, test_ids = load_data()
 
 # Validate the conditions mentioned in the docstring
-assert len(X_test) == len(test_ids), "Mismatch in length of test images and test IDs"
-assert len(X) == len(y), "Mismatch in length of training images and labels"
+assert X_test.shape[0] == test_ids.shape[0], "Mismatch in length of test images and test IDs"
+assert X.shape[0] == y.shape[0], "Mismatch in length of training images and labels"
 
 print("Data loader test passed successfully. Length of test images matches length of test IDs.")