class_docstr: The PredictUnit is an interface that can be used to organize your prediction logic. The core of it is the ``predict_step`` which
  is an abstract method where you can define the code you want to run each iteration of the dataloader.
  
  To use the PredictUnit, create a class which subclasses :class:`~torchtnt.framework.unit.PredictUnit`.
  Then implement the ``predict_step`` method on your class, and then you can 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_predict_batch`` to modify the default batch fetching behavior.
  Below is a simple example of a user's subclass of :class:`~torchtnt.framework.unit.PredictUnit` that implements a basic ``predict_step``.
  
  .. code-block:: python
  
    from torchtnt.framework.unit import PredictUnit
  
    Batch = Tuple[torch.tensor, torch.tensor]
    # specify type of the data in each batch of the dataloader to allow for typechecking
  
    class MyPredictUnit(PredictUnit[Batch]):
        def __init__(
            self,
            module: torch.nn.Module,
        ):
            super().__init__()
            self.module = module
  
        def predict_step(self, state: State, data: Batch) -> torch.tensor:
            inputs, targets = data
            outputs = self.module(inputs)
            return outputs
  
    predict_unit = MyPredictUnit(module=...)
class_docstr_tok: ["The","PredictUnit","is","an","interface","that","can","be","used","to","organize","your","prediction","logic",".","The","core","of","it","is","the","``","predict_step","``","which","is","an","abstract","method","where","you","can","define","the","code","you","want","to","run","each","iteration","of","the","dataloader",".","To","use","the","PredictUnit",",","create","a","class","which","subclasses",":","class",":","`","~torchtnt.framework.unit.PredictUnit","`",".","Then","implement","the","``","predict_step","``","method","on","your","class",",","and","then","you","can","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_predict_batch","``","to","modify","the","default","batch","fetching","behavior",".","Below","is","a","simple","example","of","a","user","'s","subclass","of",":","class",":","`","~torchtnt.framework.unit.PredictUnit","`","that","implements","a","basic","``","predict_step","``",".","..","code-block",":",":","python","from","torchtnt.framework.unit","import","PredictUnit","Batch","=","Tuple","[","torch.tensor",",","torch.tensor","]","#","specify","type","of","the","data","in","each","batch","of","the","dataloader","to","allow","for","typechecking","class","MyPredictUnit","(","PredictUnit","[","Batch","]",")",":","def","__init__","(","self",",","module",":","torch.nn.Module",",",")",":","super","(",")",".__init__","(",")","self.module","=","module","def","predict_step","(","self",",","state",":","State",",","data",":","Batch",")","-",">","torch.tensor",":","inputs",",","targets","=","data","outputs","=","self.module","(","inputs",")","return","outputs","predict_unit","=","MyPredictUnit","(","module=","...",")"]
docstr: Hook called after a predict epoch ends.
  
  Args:
      state: a :class:`~torchtnt.framework.state.State` object containing metadata about the prediction run.
docstr_tok: ["Hook","called","after","a","predict","epoch","ends",".","Args",":","state",":","a",":","class",":","`","~torchtnt.framework.state.State","`","object","containing","metadata","about","the","prediction","run","."]
code: def on_predict_epoch_end(self, state: State) -> None:
      """Hook called after a predict epoch ends.
  
      Args:
          state: a :class:`~torchtnt.framework.state.State` object containing metadata about the prediction 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 5 for calling the PredictUnit obj's underlying member method code with example usage: obj.on_predict_epoch_end(state) without return types