Update transform function for hip

ROCm · Apr 22, 2024 · 7835282 · 7835282
1 parent aedfa8f
commit 7835282
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/bitsandbytes/backends/cuda.py b/bitsandbytes/backends/cuda.py
@@ -116,6 +116,9 @@ def double_quant(self, A, col_stats=None, row_stats=None, out_col=None, out_row=
         return out_row, out_col, row_stats, col_stats, coo_tensor
 
     def transform(self, A, to_order, from_order="row", out=None, transpose=False, state=None, ld=None):
+        if HIP_ENVIRONMENT:
+            return nvidia_transform(A, to_order, from_order, out, transpose, state, ld)
+
         prev_device = pre_call(A.device)
         if state is None:
             state = (A.shape, from_order)