Skip to content

Commit

Permalink
fixed stage name
Browse files Browse the repository at this point in the history
  • Loading branch information
kozlov721 committed Nov 28, 2024
1 parent b4180c3 commit 5be4b43
Showing 1 changed file with 2 additions and 3 deletions.
5 changes: 2 additions & 3 deletions modelconverter/utils/config.py
Original file line number Diff line number Diff line change
Expand Up @@ -447,9 +447,8 @@ def _validate_stages(cls, data: Dict[str, Any]) -> Dict[str, Any]:
else:
extra = {}
for key in list(data.keys()):
if key not in cls.__fields__:
if key not in cls.model_fields:
extra[key] = data.pop(key)

for stage in data["stages"].values():
for key, value in extra.items():
if key not in stage:
Expand All @@ -459,7 +458,7 @@ def _validate_stages(cls, data: Dict[str, Any]) -> Dict[str, Any]:
@model_validator(mode="after")
def _validate_single_stage_name(self) -> Self:
"""Changes the default 'default_stage' name to the name of the input model."""
if len(self.stages) == 1:
if len(self.stages) == 1 and "default_stage" in self.stages:
stage = next(iter(self.stages.values()))
model_name = stage.input_model.stem
self.stages = {model_name: stage}
Expand Down

0 comments on commit 5be4b43

Please sign in to comment.