2399bacca070189e2f89c6a6e1dbf72239d8129f,tf_agents/policies/policy_saver_test.py,PolicySaverTest,testTrainStepSaved,#PolicySaverTest#,429

Before Change

      saver = policy_saver.PolicySaver(
          policy, batch_size=None, train_step=train_step)
      self.assertEqual(7, saver.get_train_step())
      path = os.path.join(self.get_temp_dir(), "save_model")
      saver.save(path)

      reloaded = tf.compat.v2.saved_model.load(path)

After Change

      saver = policy_saver.PolicySaver(
          policy, batch_size=None, train_step=train_step)
      if tf.executing_eagerly():
        step = saver.get_train_step()
      else:
        step = self.evaluate(saver.get_train_step())
      self.assertEqual(7, step)
      path = os.path.join(self.get_temp_dir(), "save_model")
      saver.save(path)

Instances