Enable torch.get_autocast_gpu_dtype in Dynamo tracing

fduwjj · fduwjj · commit 709eee7c3271 · 2023-06-07T18:01:12.000Z
ghstack-source-id: c343dcf Pull Request resolved: #103166
diff --git a/test/dynamo/test_functions.py b/test/dynamo/test_functions.py
@@ -399,6 +399,11 @@ def test_get_default_dtype(x):
         else:
             return x - 1
 
+    @make_test
+    def test_get_autocast_gpu_dtype(x):
+        dtype = torch.get_autocast_gpu_dtype()
+        return x.type(dtype)
+
     @make_test
     def test_is_complex(x):
         if torch.is_complex(x):
diff --git a/torch/_dynamo/variables/torch.py b/torch/_dynamo/variables/torch.py
@@ -64,6 +64,7 @@
     torch.device,
     torch.distributed.is_available,
     torch.finfo,
+    torch.get_autocast_gpu_dtype,
     torch.get_default_dtype,
     torch.iinfo,
     torch.is_autocast_cache_enabled,
