prevent divide by zero in cuda kernel (#471)

EricLBuehler · Jun 24, 2024 · 333ce88 · 333ce88
1 parent 7558a12
commit 333ce88
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/mistralrs-core/src/cuda/nonzero_bitwise.cu b/mistralrs-core/src/cuda/nonzero_bitwise.cu
@@ -66,7 +66,11 @@ __global__ void transform_indices(const uint32_t *temp_indices,
     int temp_index = temp_indices[idx];
     for (int i = num_dims - 1; i >= 0; i--) {
       d_out[idx * num_dims + i] = temp_index % dims[i];
-      temp_index /= dims[i];
+      if (dims[i] == 0) {
+        // Handle the error appropriately
+      } else {
+        temp_index /= dims[i];
+      }
     }
   }
 }