Merge pull request #81 from Gabriel2409/bugfix/distributedjob-type-hint

Make Framework an Enum
getindata · Oct 11, 2023 · 8e5979f · 8e5979f
2 parents 1a7f8dd + 8c6e139
commit 8e5979f
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/kedro_azureml/distributed/config.py b/kedro_azureml/distributed/config.py
@@ -1,9 +1,11 @@
 import json
 from dataclasses import asdict, dataclass
+from enum import Enum
 from typing import Optional, Union
 
 
-class Framework:
+# can be replaced by StrEnum starting from python 3.11
+class Framework(str, Enum):
     PyTorch = "PyTorch"
     TensorFlow = "TensorFlow"
     MPI = "MPI"