import torch
import torch.nn as nn


class Model(nn.Module):
    def __init__(self, center, scale, delta):
        super(Model, self).__init__()
        self.center = center
        self.scale = scale
        self.delta = delta

    def forward(self, x):
        y = (x - self.center) / self.scale
        abs_y = torch.abs(y)

        quad = 0.5 * y * y
        linear = self.delta * (abs_y - 0.5 * self.delta)

        return torch.where(abs_y <= self.delta, quad, linear)


batch_size = 1024
dim = 1024


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


def get_init_inputs():
    return [0.0, 1.0, 1.0]