import torch
import torch.nn as nn


class Model(nn.Module):
    def __init__(self, min_val, max_val, eps=1e-5):
        super(Model, self).__init__()
        self.min_val = min_val
        self.max_val = max_val
        self.eps = eps

    def forward(self, x):
        mean = x.mean(dim=-1, keepdim=True)
        std = x.std(dim=-1, keepdim=True)

        norm = (x - mean) / (std + self.eps)
        return torch.clamp(norm, self.min_val, self.max_val)


batch_size = 16
dim = 256
min_val = -1.0
max_val = 1.0


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


def get_init_inputs():
    return [min_val, max_val]