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


class Model(nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, y_pred: torch.Tensor, y_true: torch.Tensor) -> torch.Tensor:
        margin = 1.0 - y_true * y_pred

        loss_elementwise = F.relu(margin).pow(2)

        return loss_elementwise.mean()


batch_size = 128
feature_dim = 512


def get_inputs():
    y_pred = torch.randn(batch_size, feature_dim, dtype=torch.float32)
    y_true = torch.randint(0, 2, (batch_size, feature_dim), dtype=torch.float32) * 2.0 - 1.0
    return [y_pred, y_true]


def get_init_inputs():
    return []