code: def requireprobs(fn):
      """Will convert logits to probs if needed"""
  
      def wrapper(self, probabilities):
          # Expected shape : [n_sample, n_classes, ..., n_iterations]
          bounded = (
              torch.min(probabilities) < 0
              or torch.max(probabilities) > 1.0
          )
          if bounded or not probabilities.sum(1).allclose(1):
              probabilities = F.softmax(probabilities, 1)
          return fn(self, probabilities)
  
      return wrapper
  
file_path: baal/baal/active/heuristics/heuristics_gpu.py
imports: from typing import Callable, Optional, Sequence
  import numpy
  import torch
  import torch.nn.functional
  from baal.active.dataset.base import Dataset
  from baal import ModelWrapper
label_desc: Use image node_id 2 for calling a global function with example usage: requireprobs(fn) and returns: wrapper, fn