From ae892f577f7360cfbaef16e50fac4856123261f5 Mon Sep 17 00:00:00 2001
From: Shashank Rajput <shashank.rajput@databricks.com>
Date: Sat, 3 Feb 2024 07:13:45 +0000
Subject: [PATCH] added step to other loggers

---
 composer/loggers/cometml_logger.py   | 3 ++-
 composer/loggers/console_logger.py   | 3 ++-
 composer/loggers/file_logger.py      | 3 ++-
 composer/loggers/in_memory_logger.py | 3 ++-
 composer/loggers/mlflow_logger.py    | 3 ++-
 5 files changed, 10 insertions(+), 5 deletions(-)

diff --git a/composer/loggers/cometml_logger.py b/composer/loggers/cometml_logger.py
index 3581d862d6..6016af68a7 100644
--- a/composer/loggers/cometml_logger.py
+++ b/composer/loggers/cometml_logger.py
@@ -98,7 +98,8 @@ def init(self, state: State, logger: Logger) -> None:
             assert self.experiment is not None
             self.experiment.set_name(self.name)
 
-    def log_table(self, columns: List[str], rows: List[List[Any]], name: str = 'Table') -> None:
+    def log_table(self, columns: List[str], rows: List[List[Any]], name: str = 'Table', step: Optional[int] = None) -> None:
+        del step
         if self._enabled:
             assert self.experiment is not None
             try:
diff --git a/composer/loggers/console_logger.py b/composer/loggers/console_logger.py
index df97cdff04..6be5b1d6af 100644
--- a/composer/loggers/console_logger.py
+++ b/composer/loggers/console_logger.py
@@ -77,7 +77,8 @@ def log_hyperparameters(self, hyperparameters: Dict[str, Any]):
         # Lazy logging of hyperparameters.
         self.hparams.update(hyperparameters)
 
-    def log_table(self, columns: List[str], rows: List[List[Any]], name: str = 'Table') -> None:
+    def log_table(self, columns: List[str], rows: List[List[Any]], name: str = 'Table', step: Optional[int] = None) -> None:
+        del step
         try:
             import pandas as pd
         except ImportError as e:
diff --git a/composer/loggers/file_logger.py b/composer/loggers/file_logger.py
index 8d88a0deae..d863866760 100644
--- a/composer/loggers/file_logger.py
+++ b/composer/loggers/file_logger.py
@@ -185,7 +185,8 @@ def log_traces(self, traces: Dict[str, Any]):
                     trace_str + '\n',
                 )
 
-    def log_table(self, columns: List[str], rows: List[List[Any]], name: str = 'Table') -> None:
+    def log_table(self, columns: List[str], rows: List[List[Any]], name: str = 'Table', step: Optional[int] = None) -> None:
+        del step
         try:
             import pandas as pd
         except ImportError as e:
diff --git a/composer/loggers/in_memory_logger.py b/composer/loggers/in_memory_logger.py
index 0bf024f8c2..08225e59ff 100644
--- a/composer/loggers/in_memory_logger.py
+++ b/composer/loggers/in_memory_logger.py
@@ -72,7 +72,8 @@ def __init__(self) -> None:
     def log_hyperparameters(self, hyperparameters: Dict[str, Any]):
         self.hyperparameters.update(hyperparameters)
 
-    def log_table(self, columns: List[str], rows: List[List[Any]], name: str = 'Table') -> None:
+    def log_table(self, columns: List[str], rows: List[List[Any]], name: str = 'Table', step: Optional[int] = None) -> None:
+        del step
         try:
             import pandas as pd
         except ImportError as e:
diff --git a/composer/loggers/mlflow_logger.py b/composer/loggers/mlflow_logger.py
index e865a7ba6f..570a8c9ff9 100644
--- a/composer/loggers/mlflow_logger.py
+++ b/composer/loggers/mlflow_logger.py
@@ -182,7 +182,8 @@ def init(self, state: State, logger: Logger) -> None:
     def after_load(self, state: State, logger: Logger) -> None:
         logger.log_hyperparameters({'mlflow_experiment_id': self._experiment_id, 'mlflow_run_id': self._run_id})
 
-    def log_table(self, columns: List[str], rows: List[List[Any]], name: str = 'Table') -> None:
+    def log_table(self, columns: List[str], rows: List[List[Any]], name: str = 'Table', step: Optional[int] = None) -> None:
+        del step
         if self._enabled:
             try:
                 import pandas as pd