message: Fox for tff mnist test failure. This fixes the test case but does not solve the underlying issue of using the incorrect graph.
diff: @@ -1971,9 +1971,10 @@ class Model(Network):
  self._per_output_weighted_metrics[i], i))
  # Create a metric wrapper for each output loss.
  + if len(self.outputs) > 1:
  self._output_loss_metrics = [
  - metrics_module.SumOverBatchSize() if hasattr(loss_fn, 'reduction') else
  - metrics_module.SumOverBatchSizeMetricWrapper(loss_fn)
  + metrics_module.SumOverBatchSize() if hasattr(loss_fn, 'reduction')
  + else metrics_module.SumOverBatchSizeMetricWrapper(loss_fn)
  for loss_fn in self.loss_functions
  ]
  