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

…t is True (apache#4894)
alexwong · Feb 28, 2020 · d3e7992 · d3e7992
1 parent 0b6164d
commit d3e7992
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