Merge branch 'comet-monitor'

microsoft · Apr 26, 2024 · 99ef13d · 99ef13d
2 parents f3957ef + 42bacbf
commit 99ef13d
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/deepspeed/monitor/config.py b/deepspeed/monitor/config.py
@@ -3,7 +3,7 @@
 
 # DeepSpeed Team
 
-from typing import Optional, Literal
+from typing import Optional
 
 from deepspeed.pydantic_v1 import root_validator
 from deepspeed.runtime.config_utils import DeepSpeedConfigModel
@@ -113,9 +113,9 @@ class CometConfig(DeepSpeedConfigModel):
     Defaults to True.
     """
 
-    mode: Optional[Literal["get_or_create", "create", "get"]] = None
+    mode: Optional[str] = None
     """
-    Control how the Comet experiment is started:
+    Control how the Comet experiment is started, 3 options are possible.:
         - "get": Continue logging to an existing experiment identified by the `experiment_key` value.
         - "create": Always creates of a new experiment, useful for HPO sweeps.
         - "get_or_create" (default): Starts a fresh experiment if required, or persists logging to an existing one.

diff --git a/tests/unit/monitor/test_monitor.py b/tests/unit/monitor/test_monitor.py
@@ -110,7 +110,7 @@ def test_comet_monitor(self):
         mock_start = Mock(return_value=mock_experiment)
 
         config_dict = {
-            "train_batch_size": 1,
+            "train_batch_size": 2,
             "comet": {
                 "enabled": True,
                 "samples_log_interval": 42,