import torch
import torch.nn as nn

class Model(nn.Module):
    def __init__(self, feature_dim):
        super().__init__()
        self.register_buffer("center", torch.zeros(feature_dim))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return torch.max(torch.abs(x - self.center), dim=1).values

batch_size = 256
feature_dim = 512

def get_inputs():
    x = torch.randn(batch_size, feature_dim, dtype=torch.float32)
    return [x]

def get_init_inputs():
    return [feature_dim]