diff --git a/random_layer.py b/random_layer.py
index 713a795..5646a56 100644
--- a/random_layer.py
+++ b/random_layer.py
@@ -24,7 +24,11 @@
 from scipy.spatial.distance import cdist, pdist, squareform
 
 from sklearn.metrics import pairwise_distances
-from sklearn.utils import check_random_state, atleast2d_or_csr
+from sklearn.utils import check_random_state
+try:
+    from sklearn.utils import atleast2d_or_csr as check_array
+except ImportError:
+    from sklearn.utils import check_array
 from sklearn.utils.extmath import safe_sparse_dot
 from sklearn.base import BaseEstimator, TransformerMixin
 
@@ -108,7 +112,7 @@ def fit(self, X, y=None):
         -------
         self
         """
-        X = atleast2d_or_csr(X)
+        X = check_array(X)
 
         self._generate_components(X)
 
@@ -130,7 +134,7 @@ def transform(self, X, y=None):
         -------
         X_new : numpy array of shape [n_samples, n_components]
         """
-        X = atleast2d_or_csr(X)
+        X = check_array(X)
 
         if (self.components_ is None):
             raise ValueError('No components initialized')