title: Tensorflow & Keras: Different output shape depends on training or inferring
body: <p>I am sub-classing <code>tensorflow.keras.Model</code> to implement a certain model. Expected behavior:</p>
  
  <ol>
  <li>Training (fitting) time: returns a list of tensors including the final output and auxiliary output;</li>
  <li>Inferring (predicting) time: returns a single output tensor.</li>
  </ol>
  
  <p>And the code is:</p>
  
  <pre class="lang-py prettyprint-override"><code>class SomeModel(tensorflow.keras.Model):
      # ......
      def call(self, x, training=True):
          # ......
          return [aux1, aux2, net] if training else net
  </code></pre>
  
  <p>This is how i use it:</p>
  
  <pre class="lang-py prettyprint-override"><code>model=SomeModel(...)
  model.compile(...,
      loss=keras.losses.SparseCategoricalCrossentropy(),
      loss_weights=[0.4, 0.4, 1],...)
  # ......
  model.fit(data, [labels, labels, labels])
  </code></pre>
  
  <p>And got:</p>
  
  <blockquote>
    <p>AssertionError: in converted code:</p>
    
    <p>ipython-input-33-862e679ab098:140 call  *</p>
  
  <pre><code>   `return [aux1, aux2, net] if training else net`
  </code></pre>
    
    <p>...\tensorflow_core\python\autograph\operators\control_flow.py:918 if_stmt</p>
  </blockquote>
  
  <p>Then the problem is that the <code>if</code> statement is converted into the calculation graph and this would of course cause the problem. I found the whole stack trace is long and useless so it's not included here.</p>
  
  <p>So, is there any way to make TensorFlow generate different graph based on <code>training</code> or not?</p>