mirror of
https://github.com/hpcaitech/ColossalAI.git
synced 2025-09-02 01:28:31 +00:00
[autoparallel] add torch.nn.ReLU metainfo (#1868)
* [fx] metainfo class for auto parallel * [fx] add unit test for linear metainfo * [fx] fix bwd param for linear * [fx] modify unit test * [fx] modify unit test * [fx] modify import * [fx] modify import * [fx] modify import * [fx] move meta profiler to auto parallel * [fx] add conv metainfo class * [fx] restore profiler * [fx] restore meta profiler * [autoparallel] modify unit test * [fx] modify unit test * [autoparallel] add batchnorm metainfo class * [autoparallel] fix batchnorm unit test function declaration * [fx] restore profiler * [fx] add relu metainfo class * [fx] restore profiler * [autoparallel] modify metainfo input
This commit is contained in:
5
colossalai/auto_parallel/meta_profiler/constants.py
Normal file
5
colossalai/auto_parallel/meta_profiler/constants.py
Normal file
@@ -0,0 +1,5 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
# list of inplace operations
|
||||
INPLACE_MODULE = [nn.ReLU]
|
Reference in New Issue
Block a user