Add new fp32 ops names to amp list

apache · Jun 9, 2021 · 2dc22c6 · 2dc22c6
1 parent 577dc93
commit 2dc22c6
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/python/mxnet/amp/lists/symbol_fp16.py b/python/mxnet/amp/lists/symbol_fp16.py
@@ -615,6 +615,8 @@
     FP32_FUNCS.extend([
         '_sg_mkldnn_conv',
         '_sg_mkldnn_fully_connected',
+        '_sg_mkldnn_selfatt_qk',
+        '_sg_mkldnn_selfatt_valatt',
     ])
 
 # Functions that have to be cast to FP32 only for