optional

mosaicml · Apr 19, 2024 · 6250a44 · 6250a44
1 parent c44fafa
commit 6250a44
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/llmfoundry/utils/config_utils.py b/llmfoundry/utils/config_utils.py
@@ -236,7 +236,7 @@ def parse_source_dataset(cfg: DictConfig):
     return data_paths
 
 
-def log_dataset_uri(cfg: DictConfig) -> mlflow.data.meta_dataset.MetaDataset:
+def log_dataset_uri(cfg: DictConfig) -> Optional[mlflow.data.meta_dataset.MetaDataset]:
     """Logs dataset tracking information to MLflow."""
     if mlflow is None:
         log.warning('MLflow is not installed. Skipping dataset logging.')

diff --git a/tests/utils/test_mlflow_logging.py b/tests/utils/test_mlflow_logging.py
@@ -1,10 +1,10 @@
 # Copyright 2024 MosaicML LLM Foundry authors
 # SPDX-License-Identifier: Apache-2.0
 
+from typing import Any
 from unittest.mock import patch
 
 import pytest
-from typing import Any
 from omegaconf import OmegaConf
 
 from llmfoundry.utils.config_utils import log_dataset_uri, parse_source_dataset