fixup test

jorisvandenbossche · jorisvandenbossche · commit a5f8b65c68a7 · 2018-07-17T12:21:54.000-05:00
diff --git a/sklearn/tests/test_impute.py b/sklearn/tests/test_impute.py
@@ -7,7 +7,6 @@
 
 from sklearn.utils.testing import assert_allclose
 from sklearn.utils.testing import assert_allclose_dense_sparse
-from sklearn.utils.testing import assert_equal
 from sklearn.utils.testing import assert_array_equal
 from sklearn.utils.testing import assert_array_almost_equal
 from sklearn.utils.testing import assert_false
@@ -66,10 +65,9 @@ def test_imputation_shape():
     for strategy in ['mean', 'median', 'most_frequent', "constant"]:
         imputer = SimpleImputer(strategy=strategy)
         X_imputed = imputer.fit_transform(sparse.csr_matrix(X))
-        assert_equal(X_imputed.shape, (10, 2))
+        assert X_imputed.shape == (10, 2)
         X_imputed = imputer.fit_transform(X)
-        assert_equal(X_imputed.shape, (10, 2))
-        
+        assert X_imputed.shape == (10, 2)
 
 
 @pytest.mark.parametrize("strategy", ["const", 101, None])
