import torch.nn
import warnings

from dataclasses import dataclass, field
from typing import Type, Any, Optional, Callable
from ignite.engine import Events, Engine


default_optimizer_class = torch.optim.Adam
default_loss_class = torch.nn.MSELoss


class ModelConfig:
    def __init__(self, model_class: Type, pretrained_source: Optional[str] = None, **model_kwargs: Any):
        self.model_class = model_class
        self.pretrained_source = pretrained_source
        self.model_kwargs = model_kwargs


class OptimizerConfig:
    def __init__(self, optimizer_class: Optional[Type] = None, **optimizer_kwargs: Any):
        if optimizer_class is None:
            warnings.warn(f"Using default optimizer: {default_optimizer_class}")
            optimizer_class = default_optimizer_class

        self.optimizer_class = optimizer_class
        self.optimizer_kwargs = optimizer_kwargs


class LossConfig:
    def __init__(self, loss_class: Optional[Type] = None, **loss_kwargs: Any):
        if loss_class is None:
            warnings.warn(f"Using default loss: {default_loss_class}")
            loss_class = default_loss_class

        self.loss_class = loss_class
        self.loss_kwargs = loss_kwargs


class Callback:
    event: Events
    function: Callable[[Engine], None]


class Callbacks:
    def __init__(self, *callbacks: Callback):
        self.callbacks = callbacks

    def __iter__(self):
        return iter(self.callbacks)


@dataclass
class ModelZooEntry:
    model_config: ModelConfig
    optimizer_config: OptimizerConfig = field(default_factory=OptimizerConfig)
    loss_config: LossConfig = field(default_factory=LossConfig)
    trainer_callbacks: Callbacks = field(default_factory=Callbacks)
    evaluator_callbacks: Callbacks = field(default_factory=Callbacks)
    predictor_callbacks: Callbacks = field(default_factory=Callbacks)

6200f752a5d7bd94eb794586aa31593e7bbe9008
<mask token>


class OptimizerConfig:

    def __init__(self, optimizer_class: Optional[Type]=None, **
        optimizer_kwargs: Any):
        if optimizer_class is None:
            warnings.warn(f'Using default optimizer: {default_optimizer_class}'
                )
            optimizer_class = default_optimizer_class
        self.optimizer_class = optimizer_class
        self.optimizer_kwargs = optimizer_kwargs


class LossConfig:

    def __init__(self, loss_class: Optional[Type]=None, **loss_kwargs: Any):
        if loss_class is None:
            warnings.warn(f'Using default loss: {default_loss_class}')
            loss_class = default_loss_class
        self.loss_class = loss_class
        self.loss_kwargs = loss_kwargs


class Callback:
    event: Events
    function: Callable[[Engine], None]


class Callbacks:

    def __init__(self, *callbacks: Callback):
        self.callbacks = callbacks

    def __iter__(self):
        return iter(self.callbacks)


@dataclass
class ModelZooEntry:
    model_config: ModelConfig
    optimizer_config: OptimizerConfig = field(default_factory=OptimizerConfig)
    loss_config: LossConfig = field(default_factory=LossConfig)
    trainer_callbacks: Callbacks = field(default_factory=Callbacks)
    evaluator_callbacks: Callbacks = field(default_factory=Callbacks)
    predictor_callbacks: Callbacks = field(default_factory=Callbacks)

<mask token>


class ModelConfig:

    def __init__(self, model_class: Type, pretrained_source: Optional[str]=
        None, **model_kwargs: Any):
        self.model_class = model_class
        self.pretrained_source = pretrained_source
        self.model_kwargs = model_kwargs


class OptimizerConfig:

    def __init__(self, optimizer_class: Optional[Type]=None, **
        optimizer_kwargs: Any):
        if optimizer_class is None:
            warnings.warn(f'Using default optimizer: {default_optimizer_class}'
                )
            optimizer_class = default_optimizer_class
        self.optimizer_class = optimizer_class
        self.optimizer_kwargs = optimizer_kwargs


class LossConfig:

    def __init__(self, loss_class: Optional[Type]=None, **loss_kwargs: Any):
        if loss_class is None:
            warnings.warn(f'Using default loss: {default_loss_class}')
            loss_class = default_loss_class
        self.loss_class = loss_class
        self.loss_kwargs = loss_kwargs


class Callback:
    event: Events
    function: Callable[[Engine], None]


class Callbacks:

    def __init__(self, *callbacks: Callback):
        self.callbacks = callbacks

    def __iter__(self):
        return iter(self.callbacks)


@dataclass
class ModelZooEntry:
    model_config: ModelConfig
    optimizer_config: OptimizerConfig = field(default_factory=OptimizerConfig)
    loss_config: LossConfig = field(default_factory=LossConfig)
    trainer_callbacks: Callbacks = field(default_factory=Callbacks)
    evaluator_callbacks: Callbacks = field(default_factory=Callbacks)
    predictor_callbacks: Callbacks = field(default_factory=Callbacks)

<mask token>
default_optimizer_class = torch.optim.Adam
default_loss_class = torch.nn.MSELoss


class ModelConfig:

    def __init__(self, model_class: Type, pretrained_source: Optional[str]=
        None, **model_kwargs: Any):
        self.model_class = model_class
        self.pretrained_source = pretrained_source
        self.model_kwargs = model_kwargs


class OptimizerConfig:

    def __init__(self, optimizer_class: Optional[Type]=None, **
        optimizer_kwargs: Any):
        if optimizer_class is None:
            warnings.warn(f'Using default optimizer: {default_optimizer_class}'
                )
            optimizer_class = default_optimizer_class
        self.optimizer_class = optimizer_class
        self.optimizer_kwargs = optimizer_kwargs


class LossConfig:

    def __init__(self, loss_class: Optional[Type]=None, **loss_kwargs: Any):
        if loss_class is None:
            warnings.warn(f'Using default loss: {default_loss_class}')
            loss_class = default_loss_class
        self.loss_class = loss_class
        self.loss_kwargs = loss_kwargs


class Callback:
    event: Events
    function: Callable[[Engine], None]


class Callbacks:

    def __init__(self, *callbacks: Callback):
        self.callbacks = callbacks

    def __iter__(self):
        return iter(self.callbacks)


@dataclass
class ModelZooEntry:
    model_config: ModelConfig
    optimizer_config: OptimizerConfig = field(default_factory=OptimizerConfig)
    loss_config: LossConfig = field(default_factory=LossConfig)
    trainer_callbacks: Callbacks = field(default_factory=Callbacks)
    evaluator_callbacks: Callbacks = field(default_factory=Callbacks)
    predictor_callbacks: Callbacks = field(default_factory=Callbacks)

import torch.nn
import warnings
from dataclasses import dataclass, field
from typing import Type, Any, Optional, Callable
from ignite.engine import Events, Engine
default_optimizer_class = torch.optim.Adam
default_loss_class = torch.nn.MSELoss


class ModelConfig:

    def __init__(self, model_class: Type, pretrained_source: Optional[str]=
        None, **model_kwargs: Any):
        self.model_class = model_class
        self.pretrained_source = pretrained_source
        self.model_kwargs = model_kwargs


class OptimizerConfig:

    def __init__(self, optimizer_class: Optional[Type]=None, **
        optimizer_kwargs: Any):
        if optimizer_class is None:
            warnings.warn(f'Using default optimizer: {default_optimizer_class}'
                )
            optimizer_class = default_optimizer_class
        self.optimizer_class = optimizer_class
        self.optimizer_kwargs = optimizer_kwargs


class LossConfig:

    def __init__(self, loss_class: Optional[Type]=None, **loss_kwargs: Any):
        if loss_class is None:
            warnings.warn(f'Using default loss: {default_loss_class}')
            loss_class = default_loss_class
        self.loss_class = loss_class
        self.loss_kwargs = loss_kwargs


class Callback:
    event: Events
    function: Callable[[Engine], None]


class Callbacks:

    def __init__(self, *callbacks: Callback):
        self.callbacks = callbacks

    def __iter__(self):
        return iter(self.callbacks)


@dataclass
class ModelZooEntry:
    model_config: ModelConfig
    optimizer_config: OptimizerConfig = field(default_factory=OptimizerConfig)
    loss_config: LossConfig = field(default_factory=LossConfig)
    trainer_callbacks: Callbacks = field(default_factory=Callbacks)
    evaluator_callbacks: Callbacks = field(default_factory=Callbacks)
    predictor_callbacks: Callbacks = field(default_factory=Callbacks)

import torch.nn
import warnings

from dataclasses import dataclass, field
from typing import Type, Any, Optional, Callable
from ignite.engine import Events, Engine


default_optimizer_class = torch.optim.Adam
default_loss_class = torch.nn.MSELoss


class ModelConfig:
    def __init__(self, model_class: Type, pretrained_source: Optional[str] = None, **model_kwargs: Any):
        self.model_class = model_class
        self.pretrained_source = pretrained_source
        self.model_kwargs = model_kwargs


class OptimizerConfig:
    def __init__(self, optimizer_class: Optional[Type] = None, **optimizer_kwargs: Any):
        if optimizer_class is None:
            warnings.warn(f"Using default optimizer: {default_optimizer_class}")
            optimizer_class = default_optimizer_class

        self.optimizer_class = optimizer_class
        self.optimizer_kwargs = optimizer_kwargs


class LossConfig:
    def __init__(self, loss_class: Optional[Type] = None, **loss_kwargs: Any):
        if loss_class is None:
            warnings.warn(f"Using default loss: {default_loss_class}")
            loss_class = default_loss_class

        self.loss_class = loss_class
        self.loss_kwargs = loss_kwargs


class Callback:
    event: Events
    function: Callable[[Engine], None]


class Callbacks:
    def __init__(self, *callbacks: Callback):
        self.callbacks = callbacks

    def __iter__(self):
        return iter(self.callbacks)


@dataclass
class ModelZooEntry:
    model_config: ModelConfig
    optimizer_config: OptimizerConfig = field(default_factory=OptimizerConfig)
    loss_config: LossConfig = field(default_factory=LossConfig)
    trainer_callbacks: Callbacks = field(default_factory=Callbacks)
    evaluator_callbacks: Callbacks = field(default_factory=Callbacks)
    predictor_callbacks: Callbacks = field(default_factory=Callbacks)
