diff --git a/tests/unit/test_tma_driver.cu b/tests/unit/test_tma_driver.cu
index 3a48a7a1..d7f08c7b 100644
--- a/tests/unit/test_tma_driver.cu
+++ b/tests/unit/test_tma_driver.cu
@@ -92,7 +92,7 @@ int main() {
 
     CUtensorMap tma_desc_host;
     CUresult res = cuTensorMapEncodeTiled(
-        &tma_desc_host, CU_TENSOR_MAP_DATA_TYPE_UINT16, 2,
+        &tma_desc_host, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 2,
         d_src, gdim, gstr, tdim, tstr,
         CU_TENSOR_MAP_INTERLEAVE_NONE,
         CU_TENSOR_MAP_SWIZZLE_NONE,