Fixed initialization of defaults for autocast_dtype

Author: jkosek

model_navigator/configuration/__init__.py

@@ -708,6 +708,10 @@ class TorchConfig(CustomConfigForFormat):
     inference_mode: bool = True
     custom_args: Optional[Dict[str, Any]] = None
 
+    def __post_init__(self):
+        """Post initialization to handle correctly enums."""
+        self.autocast_dtype: AutocastType = AutocastType(self.autocast_dtype)
+
     @property
     def format(self) -> Format:
         """Returns Format.TORCH.
@@ -762,6 +766,7 @@ def __post_init__(self) -> None:
         """Parse dataclass enums."""
         jit_type = (self.jit_type,) if not isinstance(self.jit_type, (list, tuple)) else self.jit_type
         self.jit_type = tuple(JitType(j) for j in jit_type)
+        self.autocast_dtype: AutocastType = AutocastType(self.autocast_dtype)
 
     @property
     def format(self) -> Format:
@@ -800,6 +805,10 @@ class TorchExportConfig(CustomConfigForFormat):
     autocast_dtype: AutocastType = AutocastType.DEVICE
     inference_mode: bool = True
 
+    def __post_init__(self):
+        """Post initialization to handle correctly enums."""
+        self.autocast_dtype: AutocastType = AutocastType(self.autocast_dtype)
+
     @property
     def format(self) -> Format:
         """Returns Format.

tests/functional/L1_Inplace_Optimize_wo_MaxBS_Search_torch_HuggingFace_Whisper/test.sh

@@ -19,8 +19,9 @@ THIS_SCRIPT_PATH="$(realpath --relative-to="$(pwd)" "$0")"
 TEST_MODULE="$(dirname "${THIS_SCRIPT_PATH}"|sed 's/\//./g').test"
 
 
+pip install --upgrade setuptools
 pip install "transformers<=4.38.2" datasets librosa --upgrade
-pip install "numpy<=1.24.0"
+pip install "numpy<2"
 
 export NAVIGATOR_CONSOLE_OUTPUT=LOGS
 export NAVIGATOR_LOG_LEVEL=DEBUG

tests/unit/base/test_api_config.py

@@ -112,6 +112,14 @@ def test_torch_config_has_valid_name_and_format():
     assert config.custom_args is None
 
 
+def test_torch_config_handle_autocast_dtype_mapping():
+    config = TorchConfig(autocast_dtype=None)  # pytype: disable=wrong-arg-types
+    assert config.autocast_dtype is AutocastType.DEVICE
+
+    config = TorchConfig(autocast_dtype="torch.float16")  # pytype: disable=wrong-arg-types
+    assert config.autocast_dtype is AutocastType.FP16
+
+
 def test_torch_config_defaults_reset_values_to_initial():
     config = TorchConfig(
         inference_mode=False, autocast=False, autocast_dtype=AutocastType.FP16, custom_args={"key": "value"}

tests/unit/base/test_custom_configs.py

@@ -75,7 +75,7 @@ def test_default_values_for_torch_script_custom_config():
     assert config.format == Format.TORCHSCRIPT
 
 
-def test_default_values_for_torch_tensorrt_custom_confg():
+def test_default_values_for_torch_tensorrt_custom_config():
     torch_tensorrt_config = TorchTensorRTConfig()
     assert torch_tensorrt_config.precision == (
         TensorRTPrecision.FP32,