vince's fix for model step (#1329)

Unity-Technologies · Oct 15, 2018 · e5254f2 · e5254f2
1 parent 5ebcbd7
commit e5254f2
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 4 deletions.
diff --git a/ml-agents/mlagents/trainers/trainer.py b/ml-agents/mlagents/trainers/trainer.py
@@ -139,12 +139,11 @@ def update_policy(self):
         """
         raise UnityTrainerException("The update_model method was not implemented.")
 
-    def save_model(self, steps):
+    def save_model(self):
         """
         Saves the model
-        :param steps: The number of steps of training
         """
-        self.policy.save_model(steps)
+        self.policy.save_model(self.get_step)
 
     def export_model(self):
         """

diff --git a/ml-agents/mlagents/trainers/trainer_controller.py b/ml-agents/mlagents/trainers/trainer_controller.py
@@ -148,7 +148,7 @@ def _save_model(self,steps=0):
         :param saver: Tensorflow saver for session.
         """
         for brain_name in self.trainers.keys():
-            self.trainers[brain_name].save_model(steps)
+            self.trainers[brain_name].save_model()
         self.logger.info('Saved Model')
 
     def _export_graph(self):