import torch
import torch.nn as nn


class Model(nn.Module):
    def __init__(self, target):
        super(Model, self).__init__()
        self.target = nn.Parameter(target)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        dot = torch.sum(x * self.target, dim=-1)
        norm_x_sq = torch.sum(x * x, dim=-1)
        norm_target_sq = torch.sum(self.target * self.target, dim=-1)

        jaccard = dot / (norm_x_sq + norm_target_sq - dot + 1e-6)
        dice = 2.0 * jaccard / (1.0 + jaccard)
        return torch.sqrt(dice)


batch_size = 128
input_dim = 1024


def get_inputs():
    x = torch.abs(torch.randn(batch_size, input_dim))
    return [x]


def get_init_inputs():
    target = torch.abs(torch.randn(input_dim))
    return [target]