Fix passing str dtype to static cache (#33741)

Co-authored-by: Guang Yang <guangyang@fb.com>
2024-10-01 00:50:17 -07:00
parent c269c5c74d
commit 808997a634
2 changed files with 2 additions and 2 deletions
--- a/tests/utils/test_cache_utils.py
+++ b/tests/utils/test_cache_utils.py
@@ -181,7 +181,7 @@ class CacheTest(unittest.TestCase):

        set_seed(0)
        device = "cpu"
-        dtype = torch.float32
+        dtype = "bfloat16"
        cache_implementation = "static"
        attn_implementation = "sdpa"  # Export and ExecuTorch only works for SdpaAttention
        batch_size = 1