github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-08 12:30:42 +00:00
Code Issues Packages Projects Releases Wiki Activity

[workflow] fixed oom tests (#5275)

Browse Source 
* [workflow] fixed oom tests

* polish

* polish

* polish
This commit is contained in:
Frank Lee
2024-01-16 18:55:13 +08:00
committed by GitHub
parent 04244aaaf1
commit d69cd2eb89
19 changed files with 50 additions and 582 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
tests/kit/model_zoo/registry.py
View File

@@ -61,7 +61,9 @@ class ModelZooRegistry(dict):
"""
self[name] = (model_fn, data_gen_fn, output_transform_fn, loss_fn, model_attribute)
def get_sub_registry(self, keyword: Union[str, List[str]], exclude: Union[str, List[str]] = None):
def get_sub_registry(
self, keyword: Union[str, List[str]], exclude: Union[str, List[str]] = None, allow_empty: bool = False
):
"""
Get a sub registry with models that contain the keyword.
@@ -95,7 +97,8 @@ class ModelZooRegistry(dict):
if not should_exclude:
new_dict[k] = v
assert len(new_dict) > 0, f"No model found with keyword {keyword}"
if not allow_empty:
assert len(new_dict) > 0, f"No model found with keyword {keyword}"
return new_dict