# mseloss_torch.py
import torch
import torch.nn as nn
import torch.nn.functional as F


BATCH_SIZE = 16
DIM = 16384 * 16 

class Model(nn.Module):

    def forward(self, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
        # F.mse_loss 默认返回 mean reduction
        return F.mse_loss(pred, target, reduction='mean') 

def get_inputs():

    pred = torch.randn(BATCH_SIZE, DIM, dtype=torch.float32)
    target = pred + torch.rand_like(pred) * 0.1 
    return [pred, target]

def get_init_inputs():
    return []