mirror of https://github.com/hpcaitech/ColossalAI
20 lines
697 B
Python
20 lines
697 B
Python
import torch
|
|
from ..registry import meta_patched_module
|
|
|
|
|
|
@meta_patched_module.register(torch.nn.LayerNorm)
|
|
@meta_patched_module.register(torch.nn.GroupNorm)
|
|
@meta_patched_module.register(torch.nn.BatchNorm1d)
|
|
@meta_patched_module.register(torch.nn.BatchNorm2d)
|
|
@meta_patched_module.register(torch.nn.BatchNorm3d)
|
|
def torch_nn_normalize(self, input):
|
|
# check shape
|
|
if isinstance(self, torch.nn.BatchNorm1d):
|
|
assert input.dim() in [2, 3]
|
|
elif isinstance(self, torch.nn.BatchNorm2d):
|
|
assert input.dim() == 4
|
|
elif isinstance(self, torch.nn.BatchNorm3d):
|
|
assert input.dim() == 5
|
|
|
|
# normalization maintain the same shape as the input
|
|
return input.clone() |