/
githubmirror
/
ColossalAI
Обзор
Документация
Войти
/
githubmirror
/
ColossalAI
Код
Запросы
0
Пакеты
0
Релизы
0
Аналитика
Безопасность
main
tests/test_booster/test_accelerator.py
14 строк
443 B
Hongxin Liu
[misc] update pre-commit and run all files (#4752)
19 сен 2023, 09:20
Не верифицирован
19 сен 2023, 09:20
079bf3c
Код
Авторство
О чём код?
import torch.nn as nn from colossalai.booster.accelerator import Accelerator from colossalai.testing import clear_cache_before_run, parameterize @clear_cache_before_run() @parameterize("device", ["cpu", "cuda"]) def test_accelerator(device): accelerator = Accelerator(device) model = nn.Linear(8, 8) model = accelerator.configure_model(model) assert next(model.parameters()).device.type == device del model, accelerator