FIX shape calculation fix

rapidsai · Oct 6, 2023 · 58ed69d · 58ed69d
1 parent 7bb5a9f
commit 58ed69d
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/python/cuml/internals/array.py b/python/cuml/internals/array.py
@@ -234,7 +234,12 @@ def __init__(
         self._mem_type = mem_type
 
         if hasattr(data, "__cuda_array_interface__"):
+            # using CuPy allows processing delayed array wrappers
+            # like cumlarray without added complexity
             data = cp.asarray(data)
+            # need to reshape if user requests specific shape
+            if shape is not None:
+                data = data.reshape(shape)
             self._array_interface = data.__cuda_array_interface__
             if mem_type in (None, MemoryType.mirror):
                 self._mem_type = MemoryType.device