add bfloat16 test (#155)

Co-authored-by: Bowen12992 <zhangbluestars@gmail.com>
FlagOpen · Aug 7, 2024 · 2e55d66 · 2e55d66
1 parent 48942a0
commit 2e55d66
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/tests/test_distribution_ops.py b/tests/test_distribution_ops.py
@@ -20,7 +20,7 @@ def test_accuracy_normal(shape, dtype):
 
 
 @pytest.mark.parametrize("shape", DISTRIBUTION_SHAPES)
-@pytest.mark.parametrize("dtype", [torch.float16, torch.float32])
+@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
 def test_accuracy_uniform(shape, dtype):
     x = torch.randn(size=shape, dtype=dtype, device="cuda")
     with flag_gems.use_gems():

diff --git a/tests/test_tensor_constructor_ops.py b/tests/test_tensor_constructor_ops.py
@@ -12,7 +12,7 @@
 
 
 @pytest.mark.parametrize("shape", DISTRIBUTION_SHAPES)
-@pytest.mark.parametrize("dtype", [torch.float16, torch.float32])
+@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
 def test_accuracy_rand(shape, dtype):
     with flag_gems.use_gems():
         res_out = torch.rand(shape, dtype=dtype, device="cuda")
@@ -21,7 +21,7 @@ def test_accuracy_rand(shape, dtype):
 
 
 @pytest.mark.parametrize("shape", DISTRIBUTION_SHAPES)
-@pytest.mark.parametrize("dtype", [torch.float16, torch.float32])
+@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
 def test_accuracy_randn(shape, dtype):
     with flag_gems.use_gems():
         res_out = torch.randn(shape, dtype=dtype, device="cuda")