Hyperparamter metric tag not printed after job completion on ai-platform (Python 3.7, TF 2.1)

Previously when we used the same code in runtime version 1.15 with Tensorflow 1.15, those jobs got executed and increase accuracy as it is being passed a hypermetric tag for creation of ai-platform job.

Model is being executed and loss and accuracy are printed in logs but still hyperparameter tuning jobs fail.

 return tf.estimator.EstimatorSpec(
        mode=mode, loss=loss, train_op=train_op,
        training_hooks=[get_logging_hooks(loss),
                        get_logging_hooks_accuracy(accuracy[1])])
eval_metric_ops = {
      'accuracy': accuracy, 'f1_score': f1_score, 'recall': recall,
      'precision': precision}

  return tf.estimator.EstimatorSpec(mode=mode,
                                    loss=loss,
                                    eval_metric_ops=eval_metric_ops)


custom_estimator_model = tf.estimator.Estimator(
        model_fn=self.model_fn(), model_dir=model_dir,
        config=get_run_config(self.strategy))

train_spec = tf.estimator.TrainSpec(input_fn=self.input_fn,
                                    max_steps=self.train_steps)
assert (self.eval_input_fn), "Please provide eval input function"
eval_spec = tf.estimator.EvalSpec(input_fn=self.eval_input_fn,
                                  steps=self.eval_steps,
                                  exporters=self.exporters,
                                  throttle_secs=self.eval_throttle_secs)
tf.estimator.train_and_evaluate(custom_estimator_model,
                                train_spec,
                                eval_spec)

config file which is passed for hypertuning of parameters is mentioned below

config['trainingInput']['hyperparameters'] = {
            'goal': 'MAXIMIZE',
            'hyperparameterMetricTag': 'accuracy',
            'maxTrials': self.max_trials,
            'maxParallelTrials': 2,
            # 'maxFailedTrials' : 1,
            'enableTrialEarlyStopping': False,
            'params': [{
                'parameterName': 'batch-size',
                'type': 'DISCRETE',
                'discreteValues': self.batch_size}, {
                    'parameterName': 'learning-rate',
                    'type': 'DISCRETE',
                    'discreteValues': self.learning_rate}]}

When I try to launch the job with the same code hypertuning jobs are failing in tensorflow version 2.1 and python version 3.7

There are no error logs, so I didn't paste one (In logs on StackDriver its shown that the job is completed successfully).

I see in the logs that the model is being trained (epochs are running with loss and accuracy values in the logs), but after going through all the epochs all the HyperTune trials show status "Failed" and metric Tag is not being printed on console. Also EVALUATION Steps doesn't occur for hyper-parameter tuning job.