frozenleaves/pytorch
1
0
Fork 0
mirror of https://github.com/pytorch/pytorch.git synced 2025-10-20 21:14:14 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix bfloat16 autocast skip (#67822)

Browse Source 
Summary:
Per title

Pull Request resolved: https://github.com/pytorch/pytorch/pull/67822

Reviewed By: mruberry

Differential Revision: D32162605

Pulled By: ngimel

fbshipit-source-id: eb5ccf6c441231e572ec93ac8c2638d028abecad
This commit is contained in:
Natalia Gimelshein
2021-11-03 21:01:05 -07:00
committed by Facebook GitHub Bot
parent 2486061c72
commit 99c7a9f09d
1 changed files with 2 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
test/test_jit_autocast.py
View File

@ -10,6 +10,7 @@ from torch.testing._internal.common_cuda import TEST_CUDA
from torch.testing._internal.common_utils import run_tests
from torch.testing import FileCheck
TEST_BFLOAT16 = torch.cuda.is_bf16_supported()
class TestAutocast(JitTestCase):
def setUp(self):
@ -41,7 +42,7 @@ class TestAutocast(JitTestCase):
self.assertEqual(x.dtype, torch.float16)
self.assertEqual(y.dtype, torch.float32)
@unittest.skipIf(not TEST_CUDA, "No cuda")
@unittest.skipIf(not TEST_CUDA or not TEST_BFLOAT16, "No cuda bfloat16 support")
def test_linear_bf16(self):
@torch.jit.script
def fn(a, b):