From b75dfe0a1e5a3ca71ca71fc06c5347a1906f4922 Mon Sep 17 00:00:00 2001
From: Thien Tran <gau.nernst@yahoo.com.sg>
Date: Fri, 20 Sep 2024 01:12:04 +0000
Subject: [PATCH] add bf16 mapping

---
 python/cutlass/backend/epilogue.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/python/cutlass/backend/epilogue.py b/python/cutlass/backend/epilogue.py
index d39ccae3a3..7d52d53c98 100644
--- a/python/cutlass/backend/epilogue.py
+++ b/python/cutlass/backend/epilogue.py
@@ -44,6 +44,7 @@
 
 dtype2ctype = {
     DataType.f16: ctypes.c_uint16,
+    DataType.bf16: ctypes.c_uint16,
     DataType.f32: ctypes.c_float,
     DataType.f64: ctypes.c_double,
     DataType.s8: ctypes.c_int8,