import torch
import torch.nn as nn


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

    def forward(self, y_pred: torch.Tensor, y_true: torch.Tensor) -> torch.Tensor:
        diff = y_true - y_pred

        loss = torch.where(diff > 0, self.tau * diff, (self.tau - 1.0) * diff)

        return loss.mean()


batch_size = 128
feature_dim = 512


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


def get_init_inputs():
    return [0.5]