From 5c01c34793265b820f929b42cd20dd9e4ed8bc94 Mon Sep 17 00:00:00 2001 From: Oleksandr Pavlyk <21087696+oleksandr-pavlyk@users.noreply.github.com> Date: Thu, 24 Jul 2025 10:30:31 -0500 Subject: [PATCH] Fix mypy error in cutlass_gemm example --- python/examples/cutlass_gemm.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/python/examples/cutlass_gemm.py b/python/examples/cutlass_gemm.py index bba8633..1675d0c 100644 --- a/python/examples/cutlass_gemm.py +++ b/python/examples/cutlass_gemm.py @@ -34,9 +34,11 @@ def as_core_Stream(cs: nvbench.CudaStream) -> core.Stream: return core.Stream.from_handle(cs.addressof()) -def make_cp_array(arr_h: np.ndarray, dev_buf: core.Buffer, dev_id: int) -> cp.ndarray: +def make_cp_array( + arr_h: np.ndarray, dev_buf: core.Buffer, dev_id: int | None +) -> cp.ndarray: cp_memview = cp.cuda.UnownedMemory( - int(dev_buf.handle), dev_buf.size, dev_buf, dev_id + int(dev_buf.handle), dev_buf.size, dev_buf, -1 if dev_id is None else dev_id ) zero_offset = 0 return cp.ndarray(