import torch
import torch.nn as nn


class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()

    def forward(self, pred_actions: torch.Tensor, expert_actions: torch.Tensor) -> torch.Tensor:
        loss = ((pred_actions - expert_actions) ** 2).mean()
        return loss


batch_size = 32
action_dim = 4


def get_inputs():
    pred_actions = torch.randn(batch_size, action_dim)
    expert_actions = torch.randn(batch_size, action_dim)
    return [pred_actions, expert_actions]


def get_init_inputs():
    return []