[RELAY][FRONTEND][TENSORFLOW] Fix FuseBatchNorm output cast float16 e…

…rror
apache · Feb 16, 2020 · 9d1bb50 · 9d1bb50
1 parent a6c42b3
commit 9d1bb50
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