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


BATCH_SIZE = 4096
FEATURE_DIM = 1024
P_NORM = 2 

class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.criterion = nn.PairwiseDistance(p=P_NORM, eps=1e-6, keepdim=False) 

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

def get_inputs():
    x1 = torch.randn(BATCH_SIZE, FEATURE_DIM, dtype=torch.float32)
    x2 = torch.randn(BATCH_SIZE, FEATURE_DIM, dtype=torch.float32)
    return [x1, x2]

def get_init_inputs():
    return []