[RELAY][FRONTEND][TF] Fix FuseBatchNorm output cast error if need_cas…

…t is True
apache · Feb 19, 2020 · 79204cb · 79204cb
1 parent a6c42b3
commit 79204cb
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/python/tvm/relay/frontend/tensorflow.py b/python/tvm/relay/frontend/tensorflow.py
@@ -897,6 +897,7 @@ def _impl(inputs, attr, params):
                       disables=['momentum'])(inputs, attr)
 
         if need_cast:
+            out = _expr.TupleGetItem(out.astuple(), 0)
             out = _op.cast(out, dtype=attr['T'].name)
         return out
     return _impl