Merge pull request #159 from SebChw/trebacz626-patch-1

Update steps.py
SebChw · Nov 17, 2023 · ceb1f5f · ceb1f5f
2 parents 0aa4cc4 + fa6468c
commit ceb1f5f
Showing 1 changed file with 6 additions and 6 deletions.
diff --git a/art/steps.py b/art/steps.py
@@ -470,9 +470,9 @@ def get_check_stage(self):
         """Returns check stage"""
         return TrainingStage.TRAIN.value
 
-    def log_params(self):
+    def log_params(self, model):
         self.results["parameters"]["number_of_steps"] = self.number_of_steps
-        super().log_params()
+        super().log_params(model)
 
 
 class Overfit(ModelStep):
@@ -508,9 +508,9 @@ def get_check_stage(self):
         """Returns check stage"""
         return TrainingStage.TRAIN.value
 
-    def log_params(self):
+    def log_params(self, model):
         self.results["parameters"]["max_epochs"] = self.max_epochs
-        super().log_params()
+        super().log_params(model)
 
 
 class Regularize(ModelStep):
@@ -540,9 +540,9 @@ def do(self, previous_states: Dict):
         art_logger.info("Training regularized model")
         self.train(trainer_kwargs={"datamodule": self.datamodule})
 
-    def log_params(self):
+    def log_params(self, model):
         self.results["parameters"].update(self.trainer_kwargs)
-        super().log_params()
+        super().log_params(model)
 
 
 class Tune(ModelStep):