from .neurobench_model import NeuroBenchModel


class TorchModel(NeuroBenchModel):
    """The TorchModel class wraps an nn.Module."""

    def __init__(self, net):
        """
        Initializes the TorchModel class.

        Args:
            net: A PyTorch nn.Module.

        """
        super().__init__()

        self.net = net
        self.net.eval()

    def __call__(self, batch):
        """
        Wraps forward pass of torch.nn model.

        Args:
            batch: A PyTorch tensor of shape (batch, timesteps, features*)

        Returns:
            preds: either a tensor to be compared with targets or passed to
                NeuroBenchPostProcessors.

        """
        return self.net(batch)

    def __net__(self):
        """Returns the underlying network."""
        return self.net
