I am using tf.estimator.Estimator and tf.estimator.EstimatorSpec to train a model, the loss is defined as below:
tf.estimator.Estimator
tf.estimator.EstimatorSpec
logits = model(fea