import torch
import torch.nn as nn

class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()

    def forward(self, x, mean, inv_covariance):
        diff = x - mean
        temp = torch.mm(diff, inv_covariance)
        dist_sq = torch.sum(temp * diff, dim=1)
        dist = torch.sqrt(dist_sq + 1e-6)
        return dist.mean()

batch_size = 32
dim = 64

def get_inputs():
    x = torch.randn(batch_size, dim, requires_grad=True)
    mean = torch.randn(dim)
    # Generate positive definite matrix
    temp = torch.randn(dim, dim)
    inv_cov = torch.mm(temp, temp.t()) + torch.eye(dim)
    return [x, mean, inv_cov]

def get_init_inputs():
    return []