From ae2a5dcc0482c74d101463f2453eafcf8b46103f Mon Sep 17 00:00:00 2001 From: Trevor Morris Date: Thu, 4 Jun 2020 20:34:27 +0000 Subject: [PATCH] Don't add cast for TF batch norm when type isn't changing --- python/tvm/relay/frontend/tensorflow.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/tvm/relay/frontend/tensorflow.py b/python/tvm/relay/frontend/tensorflow.py index 201c6bac59f0..50987f9209f3 100644 --- a/python/tvm/relay/frontend/tensorflow.py +++ b/python/tvm/relay/frontend/tensorflow.py @@ -1227,7 +1227,7 @@ def _impl(inputs, attr, params, mod): attr['data_format'] = attr['data_format'].decode("utf-8") if attr['data_format'] == 'NCHW': axis = 1 - if 'U' in attr: + if 'U' in attr and attr['U'].name != attr['T'].name: need_cast = True inputs[0] = _op.cast(inputs[0], dtype=attr['U'].name) # Check if mean and variance are empty