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

class Model(nn.Module):
    def __init__(self, margin: float = 1.0):
        super(Model, self).__init__()
        self.margin = margin
        self.triplet_margin_loss = torch.nn.TripletMarginLoss(margin=self.margin, reduction='mean')

    def forward(self, anchor: torch.Tensor, positive: torch.Tensor, negative: torch.Tensor) -> torch.Tensor:
        return self.triplet_margin_loss(anchor, positive, negative)

# 定义标准维度
batch_size = 512
dim = 4096
margin = 1.0

def get_inputs():
    """
    为anchor, positive, 和 negative生成三个随机张量。
    """
    anchor = torch.randn(batch_size, dim)
    positive = torch.randn(batch_size, dim)
    negative = torch.randn(batch_size, dim)
    return [anchor, positive, negative]

def get_init_inputs():
    return [margin]