Autocast vdot (#125697)

Fixes #125544 Pull Request resolved: #125697 Approved by: https://github.com/jbschlosser
pytorch · May 14, 2024 · 3759676 · 3759676
1 parent 556e4ec
commit 3759676
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 0 deletions.
diff --git a/aten/src/ATen/autocast_mode.h b/aten/src/ATen/autocast_mode.h
@@ -895,6 +895,7 @@ copy pasted in from VariableTypeEverything.cpp with appropriate substitutions.
   _(bilinear)                \
   _(cross)                   \
   _(dot)                     \
+  _(vdot)                    \
   _(grid_sampler)            \
   _(index_put)               \
   _(tensordot)               \

diff --git a/torch/testing/_internal/autocast_test_lists.py b/torch/testing/_internal/autocast_test_lists.py
@@ -184,6 +184,7 @@ def __init__(self, dev):
             ("cross", (torch.randn(3, dtype=torch.float32, device=dev),
                        torch.randn(3, dtype=torch.float16, device=dev))),
             ("dot", pointwise0_fp16 + pointwise1_fp32),
+            ("vdot", pointwise0_fp16 + pointwise1_fp32),
             ("grid_sampler", (torch.randn((2, 3, 33, 22), dtype=torch.float16, device=dev),
                               torch.randn((2, 22, 11, 2), dtype=torch.float32, device=dev),
                               0, 0, False)),