class_docstr: ``StyleTransfer`` is a :class:`~flash.Task` for transferring the style from one image onto another. For more
  details, see :ref:`style_transfer`.
  
  Args:
      style_image: Image or path to an image to derive the style from.
      model: The model by the style transfer task.
      backbone: A string or model to use to compute the style loss from.
      content_layer: Which layer from the backbone to extract the content loss from.
      content_weight: The weight associated with the content loss. A lower value will lose content over style.
      style_layers: Layers from the backbone to derive the style loss from.
      style_weight: The weight associated with the style loss. A lower value will lose style over content.
      optimizer: Optimizer to use for training.
      lr_scheduler: The LR scheduler to use during training.
      learning_rate: Learning rate to use for training, defaults to ``1e-3``.
class_docstr_tok: ["``","StyleTransfer","``","is","a",":","class",":","`","~flash.Task","`","for","transferring","the","style","from","one","image","onto","another",".","For","more","details",",","see",":","ref",":","`","style_transfer","`",".","Args",":","style_image",":","Image","or","path","to","an","image","to","derive","the","style","from",".","model",":","The","model","by","the","style","transfer","task",".","backbone",":","A","string","or","model","to","use","to","compute","the","style","loss","from",".","content_layer",":","Which","layer","from","the","backbone","to","extract","the","content","loss","from",".","content_weight",":","The","weight","associated","with","the","content","loss",".","A","lower","value","will","lose","content","over","style",".","style_layers",":","Layers","from","the","backbone","to","derive","the","style","loss","from",".","style_weight",":","The","weight","associated","with","the","style","loss",".","A","lower","value","will","lose","style","over","content",".","optimizer",":","Optimizer","to","use","for","training",".","lr_scheduler",":","The","LR","scheduler","to","use","during","training",".","learning_rate",":","Learning","rate","to","use","for","training",",","defaults","to","``","1e-3","``","."]
code: def predict_step(
      self, batch: Any, batch_idx: int, dataloader_idx: int = 0
  ) -> Any:
      input_image = batch[DataKeys.INPUT]
      return self(input_image)
  
file_path: lightning-flash/src/flash/image/style_transfer/model.py
imports: from typing import Any, List, NoReturn, Optional, Sequence, Tuple, Union, cast
  from torch import Tensor, nn
  from flash.core.data.io.input import DataKeys
  from flash.core.model import Task
  from flash.core.registry import FlashRegistry
  from flash.core.utilities.imports import _TOPIC_IMAGE_AVAILABLE
  from flash.core.utilities.stability import beta
  from flash.core.utilities.types import LR_SCHEDULER_TYPE, OPTIMIZER_TYPE
  from flash.image.style_transfer import STYLE_TRANSFER_BACKBONES
  from flash.image.style_transfer.utils import raise_not_supported
label_desc: Use image node_id 8 for calling the StyleTransfer obj's underlying member method code with example usage: obj.predict_step(batch, batch_idx, dataloader_idx) and returns: self