code: def __init__(
      self,
      model: Union[nn.Module, nn.DataParallel],
      augment_fn: Union[Callable, Dict[str, Callable], List[Callable]],
      deaugment_fn: Union[
          Callable, Dict[str, Callable], List[Callable]
      ],
  ):
      super().__init__()
      self.model = model
      self.augment_fn = augment_fn
      self.deaugment_fn = deaugment_fn
  
imports: from collections import defaultdict
  from functools import partial
  from typing import Tuple, List, Optional, Union, Callable, Dict, Mapping
  import torch
  from torch import Tensor, nn
  from .None import functional
  from ..utils.support import pytorch_toolbelt_deprecated
label_desc: Use image node_id 1 to create a new GeneralizedTTA object from inherited base classes: nn with example: obj = GeneralizedTTA(model, augment_fn, deaugment_fn)