forked from ccf-ai-infra/GPUCodeForces
23 lines
713 B
Python
23 lines
713 B
Python
import torch
|
||
import torch.nn as nn
|
||
|
||
class Model(nn.Module):
|
||
def __init__(self, normalized_shape, eps=1e-5):
|
||
super(Model, self).__init__()
|
||
self.normalized_shape = normalized_shape
|
||
self.eps = eps
|
||
|
||
# 使用PyTorch内置的LayerNorm，这是标准实现
|
||
self.layer_norm = nn.LayerNorm(normalized_shape, eps=eps)
|
||
|
||
def forward(self, x):
|
||
# 使用PyTorch内置LayerNorm，这是标准的实现方式
|
||
return self.layer_norm(x)
|
||
|
||
def get_inputs():
|
||
# 使用更大的输入尺寸以获得更好的性能对比
|
||
# 增加batch size和hidden size，模拟真实场景
|
||
return [torch.randn(16, 16384)]
|
||
|
||
def get_init_inputs():
|
||
return [16384] |