Source code for mighty.utils.stub

"""
Stubs
-----

.. autosummary::
    :toctree: toctree/utils/

    OptimizerStub

"""

from typing import Optional, Callable

import torch
import torch.nn as nn
from torch.optim.optimizer import Optimizer


__all__ = [
    "OptimizerStub",
    "CriterionStub"
]




[docs]
class OptimizerStub(Optimizer):
    """
    An Optimizer stub for trainers that update model weights in a gradient-free
    fashion.
    """

    def __init__(self):
        self.param_groups = []



[docs]
    def step(self, closure: Optional[Callable[[], float]] = ...):
        pass





[docs]
    def state_dict(self) -> dict:
        return {}





[docs]
    def load_state_dict(self, state_dict: dict) -> None:
        pass






class CriterionStub(nn.Module):
    def forward(self, *args, **kwargs):
        return torch.Tensor([0.0])