Add BF16 output test coverage (#1835)

Summary: Pull Request resolved: #1835 As title Reviewed By: jspark1105 Differential Revision: D46830749 fbshipit-source-id: eb17bac05fbbc06d4dd389268a6c1e203922fca5
pytorch · Jun 19, 2023 · 9172480 · 9172480
1 parent 47691ce
commit 9172480
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,