code: def estimator_spec_train(self, loss, num_async_replicas=1, use_tpu=False):
      
      train_op = self.optimize(loss, num_async_replicas=num_async_replicas,
                               use_tpu=use_tpu)
  
      if use_tpu:
        if self._hparams.warm_start_from:
          def scaffold_fn():
            self.initialize_from_ckpt(self._hparams.warm_start_from)
            return tf.train.Scaffold()
        else:
          scaffold_fn = None
  
        # Note: important to call this before remove_summaries()
        if self.hparams.tpu_enable_host_call:
          host_call = self.create_train_host_call()
        else:
          host_call = None
  
        remove_summaries()
  
        return tf.contrib.tpu.TPUEstimatorSpec(
            tf.estimator.ModeKeys.TRAIN,
            loss=loss,
            train_op=train_op,
            host_call=host_call,
            scaffold_fn=scaffold_fn)
      else:
        if self._hparams.warm_start_from:
          self.initialize_from_ckpt(self._hparams.warm_start_from)
  
        # When loading weights from a pre-trained model, you want to be able to
        # load separate weights into the encoder and decoder.
        if self._hparams.warm_start_from_second:
          self.initialize_from_ckpt(self._hparams.warm_start_from_second)
  
        return tf.estimator.EstimatorSpec(
            tf.estimator.ModeKeys.TRAIN, loss=loss, train_op=train_op)
docs: Constructs `tf.estimator.EstimatorSpec` for TRAIN (training) mode.