mirror of
https://github.com/hpcaitech/ColossalAI.git
synced 2025-09-08 12:30:42 +00:00
[shardformer] refactored the shardformer layer structure (#4053)
This commit is contained in:
@@ -4,9 +4,11 @@ import torch
|
||||
import torch.nn as nn
|
||||
from torch.distributed import ProcessGroup
|
||||
|
||||
from .parallelmodule import ParallelModule
|
||||
from .parallel_module import ParallelModule
|
||||
from .utils import create_randomizer_with_offset
|
||||
|
||||
__all__ = ['Dropout1D']
|
||||
|
||||
|
||||
class Dropout1D(ParallelModule, nn.Dropout):
|
||||
"""
|
||||
|
Reference in New Issue
Block a user