Fix train test split pytest comparison (#4062)

Unblocks CI

Authors:
  - Dante Gama Dessavre (https://github.com/dantegd)

Approvers:
  - William Hicks (https://github.com/wphicks)
  - Corey J. Nolet (https://github.com/cjnolet)

URL: #4062

python/cuml/test/test_train_test_split.py

@@ -386,8 +386,8 @@ def test_stratify_retain_index(test_size, train_size):
                                                         shuffle=True,
                                                         stratify=y,
                                                         random_state=15)
-    assert (X_train["x"] == X_train.index).all()
-    assert (X_test["x"] == X_test.index).all()
+    assert (X_train["x"].to_array() == X_train.index.to_array()).all()
+    assert (X_test["x"].to_array() == X_test.index.to_array()).all()
 
     if train_size is not None:
         assert X_train.shape[0] == (int)(X.shape[0] * train_size)
@@ -440,8 +440,8 @@ def test_stratify_any_input(test_size, train_size):
                                                         shuffle=True,
                                                         stratify=X['test_col'],
                                                         random_state=15)
-    assert (X_train["x"] == X_train.index).all()
-    assert (X_test["x"] == X_test.index).all()
+    assert (X_train["x"].to_array() == X_train.index.to_array()).all()
+    assert (X_test["x"].to_array() == X_test.index.to_array()).all()
 
     if train_size is not None:
         assert X_train.shape[0] == (int)(X.shape[0] * train_size)