fix create_pyspark_dataframe to get it works for cp arrays as input

NVIDIA · Oct 16, 2024 · db2a2e9 · db2a2e9
1 parent d37c32d
commit db2a2e9
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/python/tests/utils.py b/python/tests/utils.py
@@ -86,6 +86,12 @@ def create_pyspark_dataframe(
     """Construct a dataframe based on features and label data."""
     assert feature_type in pyspark_supported_feature_types
 
+    # in case cp.ndarray get passed in
+    if not isinstance(data, np.ndarray):
+        data = data.get()
+    if label and not isinstance(label, np.ndarray):
+        label = label.get()
+
     m, n = data.shape
 
     pyspark_type = dtype_to_pyspark_type(dtype)