# torchcode.py (LayerNorm 版本)

import torch
import torch.nn as nn

class Model(torch.nn.Module):
    def __init__(self):
        super().__init__()
        # 创建一个 LayerNorm 层
        # normalized_shape 应该是除了 batch 维度之外的所有维度
        self.ln = nn.LayerNorm(normalized_shape=[16384])

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.ln(x)

batch_size = 16
dim = 16384

def get_inputs():
    x = torch.randn(batch_size, dim)
    return [x]

def get_init_inputs():
    return []
