from typing import Mapping, Any
import torch.nn as nn
from catalyst.dl.runner import SupervisedRunner, SupervisedWandbRunner
from catalyst.dl.core import RunnerState
from catalyst.contrib.optimizers import Lookahead


class ModelRunner(SupervisedWandbRunner):
    def __init__(
            self,
            model: nn.Module = None,
            device=None,
            # input_key: str = ("images", "meta"),
            input_key: str = "images",
            output_key: str = "logits",
            input_target_key: str = "targets",
    ):
        super(ModelRunner, self).__init__(
            model=model,
            device=device,
            input_key=input_key,
            output_key=output_key,
            input_target_key=input_target_key
        )