Add BF16 output test coverage

Summary: As title Differential Revision: D46830749 fbshipit-source-id: 5e5fe3c06e51830af98ea1509db376c41e00ece8
pytorch · Jun 18, 2023 · 0ca4e8b · 0ca4e8b
1 parent 47691ce
commit 0ca4e8b
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/fbgemm_gpu/test/split_table_batched_embeddings_test.py b/fbgemm_gpu/test/split_table_batched_embeddings_test.py
@@ -4135,7 +4135,9 @@ def test_nbit_forward_gpu_no_cache(
         else:
             weights_ty: SparseType = nbit_weights_ty
             mixed_weights_ty = False
-        output_dtype = random.choice([SparseType.FP32, SparseType.FP16])
+        output_dtype = random.choice(
+            [SparseType.FP32, SparseType.FP16, SparseType.BF16]
+        )
         self.execute_nbit_forward_(
             T,
             D,