mirror of
https://github.com/hpcaitech/ColossalAI.git
synced 2025-05-03 14:08:11 +00:00
206280408a
* [elixir] add elixir * [elixir] add unit tests * remove useless code * fix python 3.8 issue * fix typo * add test skip * add docstrings * add docstrings * add readme * fix typo
24 lines
576 B
Python
24 lines
576 B
Python
import pytest
|
|
import torch
|
|
|
|
from colossalai.testing import run_on_environment_flag
|
|
from tests.test_elixir.utils import to_cuda
|
|
|
|
|
|
@run_on_environment_flag('ELX')
|
|
def test_registry():
|
|
from tests.test_elixir.utils.registry import TEST_MODELS
|
|
for name, model_tuple in TEST_MODELS:
|
|
torch.cuda.synchronize()
|
|
print(f'model `{name}` is in testing')
|
|
|
|
model_fn, data_fn = model_tuple
|
|
model = model_fn().cuda()
|
|
data = to_cuda(data_fn())
|
|
loss = model(**data)
|
|
loss.backward()
|
|
|
|
|
|
if __name__ == '__main__':
|
|
test_registry()
|