fix

transformer_engine/pytorch/cpp_extensions/gemm.py

@@ -581,7 +581,7 @@ def general_grouped_gemm(
                 workspaces[0].shape[0],
                 accumulate,
                 use_split_accumulator,
-            )
+            )[0]
             for i in range(num_gemms):
                 out[i].copy_(dw[i])
             return out, bias, gelu_input