/
githubmirror
/
ColossalAI
Обзор
Документация
Войти
/
githubmirror
/
ColossalAI
Код
Запросы
0
Пакеты
0
Релизы
0
Аналитика
Безопасность
main
colossalai/auto_parallel/passes/constants.py
13 строк
417 B
YuliangLiu0306
[autoparallel] refactor runtime pass (#2644)
15 фев 2023, 05:36
Не верифицирован
15 фев 2023, 05:36
cb2c6a2
Код
Авторство
О чём код?
import torch OUTPUT_SAVED_OPS = [torch.nn.functional.relu, torch.nn.functional.softmax, torch.flatten] OUTPUT_SAVED_MOD = [ torch.nn.ReLU, torch.nn.Softmax, ] # SHAPE_ARGUMENT_OPS contains node with (input, *shape) style args. # This list could be extended if any other method has the same # argument style as view and reshape. SHAPE_ARGUMENT_OPS = [torch.Tensor.view, torch.Tensor.reshape, torch.reshape]