SUMIN/HuggingFace_transformer
1
0
Fork 0
Code Issues Pull Requests Actions 7 Packages Projects Releases Wiki Activity

[tests] make TestDeepSpeedModelZoo device-agnostic (#31402)

Browse Source 
* fix

* use accelerator device count

* ci fix
This commit is contained in:
Fanli Lin
2024-06-17 22:42:57 +08:00
committed by GitHub
parent 7977f206dc
commit 9454f437b0
2 changed files with 9 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
src/transformers/testing_utils.py
View File

@@ -2432,6 +2432,10 @@ if is_torch_available():
BACKEND_MANUAL_SEED = {"cuda": torch.cuda.manual_seed, "cpu": torch.manual_seed, "default": torch.manual_seed}
BACKEND_EMPTY_CACHE = {"cuda": torch.cuda.empty_cache, "cpu": None, "default": None}
BACKEND_DEVICE_COUNT = {"cuda": torch.cuda.device_count, "cpu": lambda: 0, "default": lambda: 1}
else:
BACKEND_MANUAL_SEED = {"default": None}
BACKEND_EMPTY_CACHE = {"default": None}
BACKEND_DEVICE_COUNT = {"default": lambda: 0}
def backend_manual_seed(device: str, seed: int):