import torch
import torch.nn as nn

N, C, H, W = 32, 64, 56, 56
EPS = 1e-8


class Variance(nn.Module):

    def __init__(self, unbiased=True):
        super().__init__()

        self.unbiased = unbiased

    def forward(self, x):

        x_flat = x.view(x.size(0), -1)

        D = x_flat.size(1)

        x_mean = x_flat.mean(dim=1, keepdim=True)

        x_centered = x_flat - x_mean

        ssd_sum = (x_centered ** 2).sum(dim=1)

        if self.unbiased:

            divisor = D - 1
        else:

            divisor = D

        if divisor <= 0:
            return torch.zeros_like(ssd_sum)

        variance = ssd_sum / divisor

        return torch.clamp(variance, min=0)


class Model(nn.Module):
    def __init__(self, unbiased=True):
        super().__init__()
        self.op = Variance(unbiased=unbiased)

    def forward(self, x):
        return self.op(x)


def get_inputs():
    x = torch.randn(N, C, H, W, dtype=torch.float32)
    return [x]


def get_init_inputs():
    return []