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


BATCH_SIZE = 16
EMBEDDING_DIM = 256  
DIM = BATCH_SIZE 

MARGIN = 0.5 

class Model(nn.Module):
    def forward(self, x1: torch.Tensor, x2: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
        return F.cosine_embedding_loss(x1, x2, y, margin=MARGIN, reduction='mean') 

def get_inputs():
   
    x1 = torch.randn(BATCH_SIZE, EMBEDDING_DIM, dtype=torch.float32)
    x2 = torch.randn(BATCH_SIZE, EMBEDDING_DIM, dtype=torch.float32)

    y = torch.randint(0, 2, size=(BATCH_SIZE,), dtype=torch.float32)
    y[y == 0] = -1.0 
    
    return [x1, x2, y]

def get_init_inputs():
    return []