import torch
import torch.nn as nn


class Model(nn.Module):
    def __init__(self, alpha=1.0, t=1.5):
        super().__init__()
        self.alpha = alpha
        self.t = t

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return torch.where(
            x > 0,
            x,
            self.alpha * (torch.pow(1 + (1 - self.t) * x, 1 / (1 - self.t)) - 1)
        )


batch_size = 1024
feature_dim = 1024


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


def get_init_inputs():
    return [1.0, 1.5]