# marginrankingloss_torch.py
import torch
import torch.nn as nn
import torch.nn.functional as F


BATCH_SIZE = 4096
FEATURE_DIM = 512
MARGIN = 1.0 

class Model(nn.Module):
   
    def __init__(self):
        super().__init__()
        self.criterion = nn.MarginRankingLoss(margin=MARGIN, reduction='mean') 

    def forward(self, x1: torch.Tensor, x2: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
   
        return self.criterion(x1, x2, target)

def get_inputs():

    x1 = torch.randn(BATCH_SIZE, FEATURE_DIM, dtype=torch.float32)
    x2 = torch.randn(BATCH_SIZE, FEATURE_DIM, dtype=torch.float32)
    target = torch.randint(0, 2, (BATCH_SIZE, FEATURE_DIM), dtype=torch.float32)
    target[target == 0] = -1 
    return [x1, x2, target]

def get_init_inputs():
    return []