# Copyright 2022 MosaicML Composer authors
# SPDX-License-Identifier: Apache-2.0
"""Reference for common types used throughout the composer library.
Attributes:
Batch (Any): Alias to type Any.
A batch of data can be represented in several formats, depending on the application.
PyTorchScheduler (torch.optim.lr_scheduler._LRScheduler): Alias for base class of learning rate schedulers such
as :class:`torch.optim.lr_scheduler.ConstantLR`.
JSON (str | float | int | None | List['JSON'] | Dict[str, 'JSON']): JSON Data.
Dataset (torch.utils.data.Dataset[Batch]): Alias for :class:`torch.utils.data.Dataset`.
"""
from __future__ import annotations
from typing import Any, Dict, List, Union
import torch
import torch.utils.data
from composer.utils import StringEnum
__all__ = ['Batch', 'PyTorchScheduler', 'JSON', 'MemoryFormat', 'TrainerMode', 'BreakEpochException']
Batch = Any
Dataset = torch.utils.data.Dataset[Batch]
PyTorchScheduler = torch.optim.lr_scheduler._LRScheduler
JSON = Union[str, float, int, None, List['JSON'], Dict[str, 'JSON']]
[docs]class BreakEpochException(Exception):
"""Raising this exception will immediately end the current epoch.
If you're wondering whether you should use this, the answer is no.
"""
pass
[docs]class TrainerMode(StringEnum):
"""Enum to represent which mode the Trainer is in.
Attributes:
TRAIN: In training mode.
EVAL: In evaluation mode.
PREDICT: In predict mode.
"""
TRAIN = 'train'
EVAL = 'eval'
PREDICT = 'predict'