Skip to content

GitLab

Unverified Commit f670152b authored by Yoshitomo Matsubara's avatar Yoshitomo Matsubara Committed by GitHub
Browse files

check if value exists without iteration (#5225)

parent 03d11338
Hide whitespace changes
Inline Side-by-side
Showing with 2 additions and 3 deletions
torchvision/prototype/models/_api.py
View file @ f670152b
...@@ -60,9 +60,8 @@ class WeightsEnum(Enum): ...@@ -60,9 +60,8 @@ class WeightsEnum(Enum):
@classmethod @classmethod
def from_str(cls, value: str) -> "WeightsEnum": def from_str(cls, value: str) -> "WeightsEnum":
for k, v in cls.__members__.items(): if value in cls.__members__:
if k == value: return cls.__members__[value]
return v
raise ValueError(f"Invalid value {value} for enum {cls.__name__}.") raise ValueError(f"Invalid value {value} for enum {cls.__name__}.")
def get_state_dict(self, progress: bool) -> OrderedDict: def get_state_dict(self, progress: bool) -> OrderedDict:
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment