tony/tensorneat-mend
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

add remove in state

Browse Source
This commit is contained in:
wls2002
2025-02-25 10:50:50 +08:00
parent 6d93930c98
commit b30cbdc669
1 changed files with 7 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
src/tensorneat/common/state.py
View File

@@ -23,6 +23,13 @@ class State:
raise ValueError(f"Key {key} does not exist in state") raise ValueError(f"Key {key} does not exist in state")
return State(**{**self.state_dict, **kwargs}) return State(**{**self.state_dict, **kwargs})
def remove(self, *keys):
for key in keys:
if key not in self.registered_keys():
raise ValueError(f"Key {key} does not exist in state")
return State(**{k: v for k, v in self.state_dict.items() if k not in keys})
def __getattr__(self, name): def __getattr__(self, name):
return self.state_dict[name] return self.state_dict[name]