import torch
import torch.nn as nn
import torch.nn.functional as F


class Model(nn.Module):
    def __init__(self, lambd=1.0507, alpha=1.67326):
        super().__init__()
        self.lambd = lambd
        self.alpha = alpha

    def forward(self, x: torch.Tensor) -> torch.Tensor:

        return torch.where(
            x >= 0,
            self.lambd * x,
            self.lambd * self.alpha * torch.expm1(x)
        )


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.0507, 1.67326]