mirror of
https://github.com/hpcaitech/ColossalAI.git
synced 2025-08-15 22:53:12 +00:00
9 lines
169 B
Python
9 lines
169 B
Python
import torch
|
|
|
|
OUTPUT_SAVED_OPS = [torch.nn.functional.relu, torch.nn.functional.softmax, torch.flatten]
|
|
|
|
OUTPUT_SAVED_MOD = [
|
|
torch.nn.ReLU,
|
|
torch.nn.Softmax,
|
|
]
|