isort

mosaicml · Aug 15, 2024 · 11eab11 · 11eab11
1 parent 72d100b
commit 11eab11
Showing 1 changed file with 8 additions and 3 deletions.
diff --git a/llmfoundry/callbacks/env_logging_callback.py b/llmfoundry/callbacks/env_logging_callback.py
@@ -13,6 +13,7 @@
 from composer.core import Callback, State
 from composer.loggers import Logger
 from composer.utils import dist
+
 from mcli import sdk
 
 __all__ = ['EnvironmentLoggingCallback']
@@ -29,6 +30,7 @@
     'peft',
 ]
 
+
 class EnvironmentLoggingCallback(Callback):
     """A callback for logging environment information during model training.
 
@@ -118,8 +120,10 @@ def _get_nvidia_info(self) -> dict[str, Any]:
         if torch.cuda.is_available():
             nccl_version = torch.cuda.nccl.version()  # type: ignore
             return {
-                'cuda_version': torch.version.cuda,  # type: ignore[attr-defined]
-                'cudnn_version': str(torch.backends.cudnn.version()),  # type: ignore[attr-defined]
+                'cuda_version':
+                    torch.version.cuda,  # type: ignore[attr-defined]
+                'cudnn_version': str(torch.backends.cudnn.version()
+                                    ),  # type: ignore[attr-defined]
                 'nccl_version': '.'.join(map(str, nccl_version)),
             }
         return {'available': False}
@@ -144,7 +148,8 @@ def fit_start(self, state: State, logger: Logger) -> None:
         # Collect environment data
         if self.log_git:
             self.env_data['git_info'] = {
-                folder: self._get_git_info(os.path.join(self.workspace_dir, folder))
+                folder:
+                self._get_git_info(os.path.join(self.workspace_dir, folder))
                 for folder in os.listdir(self.workspace_dir)
                 if os.path.isdir(os.path.join(self.workspace_dir, folder))
             }