upd

flashinfer-ai · Sep 30, 2024 · 976791c · 976791c
1 parent cfcbd93
commit 976791c
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 3 deletions.
diff --git a/include/flashinfer/group_gemm/wrapper.cuh b/include/flashinfer/group_gemm/wrapper.cuh
@@ -299,7 +299,7 @@ cudaError_t CutlassSegmentGEMMWrapper(CutlassSegmentGEMMHandler* handler, DTypeI
         // Initialize the gemm kernel
         Gemm gemm;
 
-    // Using the arguments, query for extra workspace required for matrix multiplication computation
+        // Using the arguments, query for extra workspace required for matrix multiplication computation
         size_t workspace_size = Gemm::get_workspace_size(arguments);
 
         // Allocate workspace memory

diff --git a/python/tests/test_group_gemm.py b/python/tests/test_group_gemm.py
@@ -86,8 +86,7 @@ def test_segment_gemm(
                 .cpu()
                 .numpy(),
                 rtol=1e-3,
-                atol=1e-3,
-                err_msg="assertion failed at batch {}".format(i),
+                atol=1e-3
             )
     else:
         torch.testing.assert_close(