From a5c1e0f86fab69709f48d5a060eb927d609fa15f Mon Sep 17 00:00:00 2001 From: Ben Barsdell Date: Fri, 28 Jun 2019 16:00:18 -0700 Subject: [PATCH] Add BatchMatMulV2 to auto_mixed_precision whitelist --- tensorflow/core/grappler/optimizers/auto_mixed_precision_lists.h | 1 + 1 file changed, 1 insertion(+) diff --git a/tensorflow/core/grappler/optimizers/auto_mixed_precision_lists.h b/tensorflow/core/grappler/optimizers/auto_mixed_precision_lists.h index e093014c7dd..fa8e9fbcc98 100644 --- a/tensorflow/core/grappler/optimizers/auto_mixed_precision_lists.h +++ b/tensorflow/core/grappler/optimizers/auto_mixed_precision_lists.h @@ -63,6 +63,7 @@ class AutoMixedPrecisionLists { auto list = gtl::FlatSet { #if CUDA_VERSION >= 9010 // Fp16 BatchMatMul is slow before CUDA 9.1. "BatchMatMul", + "BatchMatMulV2", #endif "BlockLSTM", "BlockLSTMGrad", "Conv2D", "Conv2DBackpropFilter", "Conv2DBackpropInput",