text: from typing import Optional, Sequence, TYPE_CHECKING, Union
  
  from torch import nn
  
  from catalyst.core.callback import Callback, CallbackNode, CallbackOrder
  from catalyst.registry import REGISTRY
  
  if TYPE_CHECKING:
      from catalyst.core.runner import IRunner
  
  
  class BatchTransformCallback(Callback):
      """Callback to perform data augmentations on GPU using kornia library.
  
      Args:
          transform: define augmentations to apply on a batch
  
              If a sequence of transforms passed, then each element
              should be either ``kornia.augmentation.AugmentationBase2D``,
              ``kornia.augmentation.AugmentationBase3D``, or ``nn.Module``
              compatible with kornia interface.
  
              If a sequence of params (``dict``) passed, then each
              element of the sequence must contain ``'transform'`` key with
              an augmentation name as a value. Please note that in this case
              to use custom augmentation you should add it to the
              `REGISTRY` registry first.
          input_key (Union[str, int]): key in batch dict mapping to transform, e.g. `'image'`
          output_key: key to use to store the result
              of the transform, defaults to `input_key` if not provided
  
      Look at `Kornia: an Open Source Differentiable Computer Vision
      Library for PyTorch`_ for details.
  
      Usage example for notebook API:
  
      .. code-block:: python
  
          import os
  
          from kornia import augmentation
          import torch
          from torch.nn import functional as F
          from torch.utils.data import DataLoader
  
          from catalyst import dl
          from catalyst.contrib.data.transforms import ToTensor
          from catalyst.contrib.datasets import MNIST
          from catalyst.contrib.callbacks.kornia_transform import (
              BatchTransformCallback
          )
          from catalyst import metrics
  
  
          class CustomRunner(dl.Runner):
              def predict_batch(self, batch):
                  # model inference step
                  return self.model(
                      batch[0].to(self.device).view(batch[0].size(0), -1)
                  )
  
              def handle_batch(self, batch):
                  # model train/valid step
                  x, y = batch
                  y_hat = self.model(x.view(x.size(0), -1))
  
                  loss = F.cross_entropy(y_hat, y)
                  accuracy01, *_ = metrics.accuracy(y_hat, y)
                  self.batch_metrics.update(
                      {"loss": loss, "accuracy01": accuracy01}
                  )
  
                  if self.is_train_loader:
                      loss.backward()
                      self.optimizer.step()
                      self.optimizer.zero_grad()
  
          model = torch.nn.Linear(28 * 28, 10)
          optimizer = torch.optim.Adam(model.parameters(), lr=0.02)
  
          loaders = {
              "train": DataLoader(
                  MNIST(os.getcwd(), train=True, transform=ToTensor()),
                  batch_size=32,
              ),
              "valid": DataLoader(
                  MNIST(os.getcwd(), train=False, transform=ToTensor()),
                  batch_size=32,
              ),
          }
          transforms = [
              augmentation.RandomAffine(degrees=(-15, 20), scale=(0.75, 1.25)),
          ]
  
          runner = CustomRunner()
  
          # model training
          runner.train(
              model=model,
              optimizer=optimizer,
              loaders=loaders,
              logdir="./logs",
              num_epochs=5,
              verbose=True,
              callbacks=[BatchTransformCallback(transforms, input_key=0)],
          )
  
      To apply augmentations only during specific loader e.g. only during
      training :class:`catalyst.core.callbacks.control_flow.ControlFlowCallback`
      callback can be used. For config API it can look like this:
  
      .. code-block:: yaml
  
          callbacks_params:
            ...
            train_transforms:
              _wrapper:
                name: ControlFlowCallback
                loaders: train
              name: BatchTransformCallback
              transform:
                - transform: kornia.RandomAffine
                  degrees: [-15, 20]
                  scale: [0.75, 1.25]
                  return_transform: true
                - transform: kornia.ColorJitter
                  brightness: 0.1
                  contrast: 0.1
                  saturation: 0.1
                  return_transform: false
              input_key: image
            ...
  
      .. _`Kornia: an Open Source Differentiable Computer Vision Library
          for PyTorch`: https://arxiv.org/pdf/1910.02190.pdf
      """
  
      def __init__(
          self,
          transform: Sequence[Union[dict, nn.Module]],
          input_key: Union[str, int] = "image",
          output_key: Optional[Union[str, int]] = None,
      ) -> None:
          """Init."""
          super().__init__(order=CallbackOrder.Internal, node=CallbackNode.all)
  
          self.input_key = input_key
          self.output_key = output_key or self.input_key
  
          transforms: Sequence[nn.Module] = [
              item if isinstance(item, nn.Module) else REGISTRY.get_from_params(**item)
              for item in transform
          ]
          assert all(
              isinstance(t, nn.Module) for t in transforms
          ), "`nn.Module` should be a base class for transforms"
  
          self.transform = nn.Sequential(*transforms)
  
      def on_batch_start(self, runner: "IRunner") -> None:
          """Apply transforms.
  
          Args:
              runner: сurrent runner
          """
          input_batch = runner.batch[self.input_key]
          output_batch = self.transform(input_batch)
          runner.batch[self.output_key] = output_batch
  
  
  __all__ = ["BatchTransformCallback"]
  