update amp list

apache · Oct 26, 2021 · bf0b9a6 · bf0b9a6
1 parent e491a89
commit bf0b9a6
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 0 deletions.
diff --git a/python/mxnet/amp/lists/symbol_bf16.py b/python/mxnet/amp/lists/symbol_bf16.py
@@ -362,6 +362,7 @@
     'zeros_like',
     '_sg_onednn_conv',
     '_sg_onednn_fully_connected',
+    '_sg_onednn_batch_dot',
     'broadcast_mul',
     'Convolution_v1',
     'IdentityAttachKLSparseReg',

diff --git a/python/mxnet/amp/lists/symbol_fp16.py b/python/mxnet/amp/lists/symbol_fp16.py
@@ -615,6 +615,7 @@
         '_sg_onednn_fully_connected',
         '_sg_onednn_selfatt_qk',
         '_sg_onednn_selfatt_valatt',
+        '_sg_onednn_batch_dot'
     ])
 
 # Functions that have to be cast to FP32 only for