code: def select_using_loss(
      batch: Union[torch.Tensor, torch.Tensor],
      batch_idx: int,
      trainer: "pl.Trainer",
      pl_module: "pl.LightningModule",
      keep: float = 0.5,
      scale_factor: float = 1,
      loss_fn: Callable = None,
  ) -> Tuple[torch.Tensor, torch.Tensor]:
      """
      Prunes minibatches as a subroutine of :class:`.SelectiveBackprop`.
  
      Computes the loss function on the provided training examples and runs minibatches
      according to the difficulty. The fraction of the minibatch that is kept for gradient
      computation is specified by the argument ``0 <= keep <= 1``.
  
      To speed up SB's selection forward pass, the argument ``scale_factor`` can
      be used to spatially downsample input tensors. The full-sized inputs
      will still be used for the weight gradient computation.
  
      Args:
          batch (Union[torch.Tensor, torch.Tensor]): Batch to prune.
          batch_idx (int): Index of the batch.
          trainer (pl.Trainer): Current trainer, used for getting current model.
          pl_module (pl.LightningModule): Current module, used for getting loss function.
          keep (float, optional): Fraction of examples in the batch to keep. Default: ``0.5``.
          scale_factor (float, optional): Multiplier between 0 and 1 for spatial size. Downsampling
              requires the input tensor to be at least 3D. Default: ``1``.
          loss_fn (Callable): Loss function of the form
              ``loss(outputs, targets, reduction='none')``.
              The function must take the keyword argument ``reduction='none'``
              to ensure that per-sample losses are returned.
  
      Returns
      -------
          (torch.Tensor, torch.Tensor): The pruned batch of inputs and targets
  
      """
      INTERPOLATE_MODES = {3: "linear", 4: "bilinear", 5: "trilinear"}
  
      input, target = batch[0], batch[1]
  
      interp_mode = "bilinear"
  
      if scale_factor > 1:
          invalidInputError(False, "scale_factor must be <= 1")
  
      if scale_factor != 1:
          if input.dim() not in INTERPOLATE_MODES:
              invalidInputError(
                  False,
                  f"Input must be 3D, 4D, \
                  or 5D if scale_factor != 1, got {input.dim()}",
              )
          interp_mode = INTERPOLATE_MODES[input.dim()]
  
      with torch.no_grad():
          N = input.shape[0]
  
          # Maybe interpolate
          if scale_factor < 1:
              X_scaled = F.interpolate(
                  input,
                  scale_factor=scale_factor,
                  mode=interp_mode,
                  align_corners=False,
                  recompute_scale_factor=False,
              )
          else:
              X_scaled = input
  
          # Get per-examples losses
          if loss_fn is None:
              invalidInputError(
                  False,
                  "loss_fn must be passed explicitly to the class.",
              )
          else:
              losses = loss_fn(trainer.model(input), target)
  
          # Check losses' dimension
          if not len(losses) == len(target):
              invalidInputError(
                  False,
                  "Losses have wrong dimension, \
              maybe they are reduced. \
              Please offer unreduced losses which have the same dimension with batch_size. \
              It can be passed by ``loss_fn=`` when you initialize the class.",
              )
  
          # Sort losses
          sorted_idx = torch.argsort(torch.Tensor(losses))
          n_select = int(keep * N)
  
          # Sample by loss
          percs = np.arange(0.5, N, 1) / N
          probs = percs ** ((1.0 / keep) - 1.0)
          probs = probs / np.sum(probs)
          select_percs_idx = np.random.choice(
              N, n_select, replace=False, p=probs
          )
          select_idx = sorted_idx[list(select_percs_idx)]
  
      return input[select_idx], target[select_idx]
  
file_path: BigDL/python/nano/src/bigdl/nano/pytorch/algorithms/selective_backprop/selective_backprop.py
imports: from typing import Any, Callable, Tuple, Union
  import warnings
  import numpy
  from bigdl.nano.utils.common import invalidInputError
  import torch
  from torch.nn import functional
  import pytorch_lightning
  from pytorch_lightning.callbacks import Callback
label_desc: Use image node_id 2 for calling a global function with example usage: select_using_loss(batch, batch_idx, trainer, pl_module, keep, scale_factor, loss_fn) and returns: input,  target