Skip to content

GitLab

Unverified Commit c50bfb80 authored by アマデウス's avatar アマデウス Committed by GitHub
Browse files

[model checkpoint] updated saving/loading for 1d layers (#594)

parent 7636d518
Expand all Hide whitespace changes
Inline Side-by-side
Showing with 371 additions and 30 deletions
colossalai/nn/layer/parallel_1d/__init__.py
View file @ c50bfb80
from .layers import (Classifier1D, Dropout1D, Embedding1D, Linear1D, Linear1D_Col, Linear1D_Row, from .layers import (Classifier1D, Dropout1D, Embedding1D, LayerNorm1D, Linear1D, Linear1D_Col, Linear1D_Row,
VocabParallelClassifier1D, VocabParallelEmbedding1D) PatchEmbedding1D, VocabParallelClassifier1D, VocabParallelEmbedding1D)
__all__ = [ __all__ = [
'Linear1D', 'Linear1D_Col', 'Linear1D_Row', 'Embedding1D', 'Dropout1D', 'Classifier1D', 'VocabParallelClassifier1D', 'Linear1D', 'Linear1D_Col', 'Linear1D_Row', 'Embedding1D', 'Dropout1D', 'Classifier1D', 'VocabParallelClassifier1D',
'VocabParallelEmbedding1D' 'VocabParallelEmbedding1D', 'LayerNorm1D', 'PatchEmbedding1D'
] ]
colossalai/nn/layer/parallel_1d/layers.py
View file @ c50bfb80
This diff is collapsed.
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment