import torch
import torch.nn as nn


class Model(nn.Module):
    def __init__(self, eps=1e-5):
        super().__init__()
        self.eps = eps

    def forward(self, img1: torch.Tensor, img2: torch.Tensor, z1: torch.Tensor, z2: torch.Tensor) -> torch.Tensor:
        img_diff = torch.abs(img1 - img2).view(img1.size(0), -1).mean(dim=1)
        z_diff = torch.abs(z1 - z2).view(z1.size(0), -1).mean(dim=1)

        loss = z_diff / (img_diff + self.eps)
        return loss.mean()


batch_size = 32
c, h, w = 3, 64, 64
z_dim = 128


def get_inputs():
    img1 = torch.randn(batch_size, c, h, w, dtype=torch.float32)
    img2 = torch.randn(batch_size, c, h, w, dtype=torch.float32)
    z1 = torch.randn(batch_size, z_dim, dtype=torch.float32)
    z2 = torch.randn(batch_size, z_dim, dtype=torch.float32)
    return [img1, img2, z1, z2]


def get_init_inputs():
    return [1e-5]