mirror of
https://github.com/hpcaitech/ColossalAI.git
synced 2025-06-28 00:07:29 +00:00
079bf3cb26
* [misc] update pre-commit * [misc] run pre-commit * [misc] remove useless configuration files * [misc] ignore cuda for clang-format
14 lines
296 B
Python
14 lines
296 B
Python
import operator
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
|
|
# list of inplace module
|
|
INPLACE_MODULE = [nn.ReLU]
|
|
|
|
# list of inplace operations
|
|
INPLACE_OPS = [torch.flatten]
|
|
|
|
# list of operations that do not save forward activations
|
|
NO_SAVE_ACTIVATION = [torch.add, torch.sub, operator.add, operator.sub]
|