[Grappler] Add Einsum to auto_mixed_precision whitelist
- Einsum is used in some of the official NLP models, and without whitelisting it they see no speedup from auto_mixed_precision. - The op calls into blas gemm routines, and so it should be treated the same as MatMul.
This commit is contained in:
parent
fd05051846
commit
0107439506
@ -69,6 +69,7 @@ class AutoMixedPrecisionLists {
|
||||
"CudnnRNNBackpropV3",
|
||||
"CudnnRNNV2",
|
||||
"CudnnRNNV3",
|
||||
"Einsum",
|
||||
"GRUBlockCell",
|
||||
"GRUBlockCellGrad",
|
||||
"LSTMBlockCell",
|
||||
|
Loading…
Reference in New Issue
Block a user