import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Variance implementation.
    Computes the variance of input tensors along the feature dimension.
    """
    def __init__(self):
        super(Model, self).__init__()
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Compute variance of input tensor.

        Args:
            x (torch.Tensor): Input tensor [batch_size, feature_dim]

        Returns:
            torch.Tensor: Variance values [batch_size]
        """
        # Compute mean
        mean = torch.mean(x, dim=1, keepdim=True)  # [batch_size, 1]
        
        # Compute squared differences
        diff = x - mean  # [batch_size, feature_dim]
        squared_diff = torch.pow(diff, 2)  # [batch_size, feature_dim]
        
        # Compute variance
        variance = torch.mean(squared_diff, dim=1)  # [batch_size]
        
        return variance

batch_size = 256
feature_dim = 1024

def get_inputs():
    # Generate input tensor with some variance
    x = torch.randn(batch_size, feature_dim) * 2.0 + 1.0  # mean=1, std=2
    return [x]

def get_init_inputs():
    return []  # No special initialization inputs needed
