class_docstr: The TrainUnit is an interface that can be used to organize your training logic. The core of it is the ``train_step`` which
  is an abstract method where you can define the code you want to run each iteration of the dataloader.
  
  To use the TrainUnit, create a class which subclasses TrainUnit. Then implement the ``train_step`` method on your class, and optionally
  implement any of the hooks, which allow you to control the behavior of the loop at different points.
  
  In addition, you can override ``get_next_train_batch`` to modify the default batch fetching behavior.
  
  Below is a simple example of a user's subclass of TrainUnit that implements a basic ``train_step``, and the ``on_train_epoch_end`` hook.
  
  .. code-block:: python
  
    from torchtnt.framework.unit import TrainUnit
  
    Batch = Tuple[torch.tensor, torch.tensor]
    # specify type of the data in each batch of the dataloader to allow for typechecking
  
    class MyTrainUnit(TrainUnit[Batch]):
        def __init__(
            self,
            module: torch.nn.Module,
            optimizer: torch.optim.Optimizer,
            lr_scheduler: torch.optim.lr_scheduler._LRScheduler,
        ):
            super().__init__()
            self.module = module
            self.optimizer = optimizer
            self.lr_scheduler = lr_scheduler
  
        def train_step(self, state: State, data: Batch) -> None:
            inputs, targets = data
            outputs = self.module(inputs)
            loss = torch.nn.functional.binary_cross_entropy_with_logits(outputs, targets)
            loss.backward()
  
            self.optimizer.step()
            self.optimizer.zero_grad()
  
        def on_train_epoch_end(self, state: State) -> None:
            # step the learning rate scheduler
            self.lr_scheduler.step()
  
    train_unit = MyTrainUnit(module=..., optimizer=..., lr_scheduler=...)
class_docstr_tok: ["The","TrainUnit","is","an","interface","that","can","be","used","to","organize","your","training","logic",".","The","core","of","it","is","the","``","train_step","``","which","is","an","abstract","method","where","you","can","define","the","code","you","want","to","run","each","iteration","of","the","dataloader",".","To","use","the","TrainUnit",",","create","a","class","which","subclasses","TrainUnit",".","Then","implement","the","``","train_step","``","method","on","your","class",",","and","optionally","implement","any","of","the","hooks",",","which","allow","you","to","control","the","behavior","of","the","loop","at","different","points",".","In","addition",",","you","can","override","``","get_next_train_batch","``","to","modify","the","default","batch","fetching","behavior",".","Below","is","a","simple","example","of","a","user","'s","subclass","of","TrainUnit","that","implements","a","basic","``","train_step","``",",","and","the","``","on_train_epoch_end","``","hook",".","..","code-block",":",":","python","from","torchtnt.framework.unit","import","TrainUnit","Batch","=","Tuple","[","torch.tensor",",","torch.tensor","]","#","specify","type","of","the","data","in","each","batch","of","the","dataloader","to","allow","for","typechecking","class","MyTrainUnit","(","TrainUnit","[","Batch","]",")",":","def","__init__","(","self",",","module",":","torch.nn.Module",",","optimizer",":","torch.optim.Optimizer",",","lr_scheduler",":","torch.optim.lr_scheduler._LRScheduler",",",")",":","super","(",")",".__init__","(",")","self.module","=","module","self.optimizer","=","optimizer","self.lr_scheduler","=","lr_scheduler","def","train_step","(","self",",","state",":","State",",","data",":","Batch",")","-",">","None",":","inputs",",","targets","=","data","outputs","=","self.module","(","inputs",")","loss","=","torch.nn.functional.binary_cross_entropy_with_logits","(","outputs",",","targets",")","loss.backward","(",")","self.optimizer.step","(",")","self.optimizer.zero_grad","(",")","def","on_train_epoch_end","(","self",",","state",":","State",")","-",">","None",":","#","step","the","learning","rate","scheduler","self.lr_scheduler.step","(",")","train_unit","=","MyTrainUnit","(","module=","...",",","optimizer=","...",",","lr_scheduler=","...",")"]
docstr: Hook called before a train epoch starts.
  
  Args:
      state: a :class:`~torchtnt.framework.state.State` object containing metadata about the training run.
docstr_tok: ["Hook","called","before","a","train","epoch","starts",".","Args",":","state",":","a",":","class",":","`","~torchtnt.framework.state.State","`","object","containing","metadata","about","the","training","run","."]
code: def on_train_epoch_start(self, state: State) -> None:
      """Hook called before a train epoch starts.
  
      Args:
          state: a :class:`~torchtnt.framework.state.State` object containing metadata about the training run.
      """
      pass
  
imports: import logging
  from abc import ABC, abstractmethod
  from typing import Any, cast, Dict, Generic, Iterator, TypeVar, Union
  import torch
  from torchtnt.framework._loop_utils import _step_requires_iterator
  from torchtnt.framework._unit_utils import _find_optimizers_for_module
  from torchtnt.framework.state import State
  from torchtnt.utils.lr_scheduler import TLRScheduler
  from torchtnt.utils.prepare_module import _is_fsdp_module, FSDPOptimizerWrapper
  from torchtnt.utils.progress import Progress
  from torchtnt.utils.stateful import Stateful
label_desc: Use image node_id 3 for calling the TrainUnit obj's underlying member method code with example usage: obj.on_train_epoch_start(state) without return types