DummyModel.py

import torch
import torch.nn


class DummyModel(torch.nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, x):
        return torch.softmax(torch.rand(5), 0)

    def __call__(self, x):
        return self.forward(x)