import torch
import torch.nn as nn

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

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        diff = x - self.mean
        temp = torch.mm(diff, self.inv_cov)
        dist_sq = (temp * diff).sum(dim=1)
        return torch.sqrt(dist_sq)

batch_size = 512
feature_dim = 256

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

def get_init_inputs():
    return [feature_dim]