log image fix

Signed-off-by: Jesse Chan <[email protected]>
mosaicml · May 22, 2024 · 604c85d · 604c85d
1 parent 57c7b72
commit 604c85d
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/composer/loggers/mlflow_logger.py b/composer/loggers/mlflow_logger.py
@@ -482,8 +482,9 @@ def log_images(
                 assert isinstance(self._run_id, str)
                 self._mlflow_client.log_image(
                     image=image,
-                    artifact_file=f'{name}_{step}_{im_ind}.png',
+                    key=f'{name}_{step}_{im_ind}',
                     run_id=self._run_id,
+                    step=step,
                 )
 
     def post_close(self):