From 83ae20c740e82a9ee488511ed0bc24e5f5a2cd3f Mon Sep 17 00:00:00 2001
From: Bogumil Sapinski Mobica
 <48835513+Bogumil-Sapinski-Mobica@users.noreply.github.com>
Date: Wed, 23 Oct 2024 18:48:31 +0200
Subject: [PATCH] added mapping for bf16 to torch::kBFloat16 (#1843)

Co-authored-by: Haicheng Wu <57973641+hwu36@users.noreply.github.com>
---
 python/cutlass/emit/pytorch.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/python/cutlass/emit/pytorch.py b/python/cutlass/emit/pytorch.py
index f7bdfd74..8c10f87a 100644
--- a/python/cutlass/emit/pytorch.py
+++ b/python/cutlass/emit/pytorch.py
@@ -255,6 +255,7 @@ _CUTLASS_TYPE_TO_TORCH_TYPE = {
     DataType.f64: "torch::kF64",
     DataType.s8: "torch::kI8",
     DataType.s32: "torch::kI32",
+    DataType.bf16: "torch::kBFloat16",
 }
 
 _PYTORCH_GEMM_IMPL_TEMPLATE_2x = (