import torch
import torch.nn as nn


class Model(nn.Module):
    def __init__(self, reduction='mean'):
        super().__init__()
        self.reduction = reduction
        self.mse_loss = nn.MSELoss(reduction=reduction)

    def forward(self, input_features: list[torch.Tensor], target_features: list[torch.Tensor]) -> torch.Tensor:
        loss = 0.0
        for i in range(len(input_features)):
            loss += self.mse_loss(input_features[i], target_features[i])
        return loss


batch_size = 256
feature_shapes = [(64, 64, 64), (128, 32, 32), (256, 16, 16)]


def get_inputs():
    input_features = []
    target_features = []
    for c, h, w in feature_shapes:
        input_features.append(torch.randn(batch_size, c, h, w, dtype=torch.float32))
        target_features.append(torch.randn(batch_size, c, h, w, dtype=torch.float32))
    return [input_features, target_features]


def get_init_inputs():
    return ['mean']