[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",
|
"CudnnRNNBackpropV3",
|
||||||
"CudnnRNNV2",
|
"CudnnRNNV2",
|
||||||
"CudnnRNNV3",
|
"CudnnRNNV3",
|
||||||
|
"Einsum",
|
||||||
"GRUBlockCell",
|
"GRUBlockCell",
|
||||||
"GRUBlockCellGrad",
|
"GRUBlockCellGrad",
|
||||||
"LSTMBlockCell",
|
"LSTMBlockCell",
|
||||||
|
Loading…
x
Reference in New Issue
Block a user