code: def extract_info_from_torch_data(
      model: torch.nn.Module,
      dataloader: Union[DataLoader, Sequence],
      dynamic_axis: Dict,
      device: Device,
      is_diffusion: bool = False,
  ):
      from nebullvm.tools.utils import ifnone
  
      input_data = (
          dataloader[0]
          if isinstance(dataloader, Sequence)
          else next(iter(dataloader))
      )
      input_row = input_data[0]
      batch_size = int(input_row[0].shape[0])
      if not all(
          [input_row[0].shape[0] == x.shape[0] for x in input_row]
      ):
          logger.warning(
              "Detected not consistent batch size in the inputs."
          )
  
      input_sizes = [tuple(x.shape) for x in input_row]
      input_types = [
          "int64"
          if isinstance(x.cpu(), torch.LongTensor)
          else "int32"
          if isinstance(x.cpu(), torch.IntTensor)
          else "float16"
          if isinstance(x.cpu(), torch.HalfTensor)
          else "float32"
          for x in input_row
      ]
  
      # For the Stable Diffusion UNet we must provide dynamic axis
      # even when using static shapes, because otherwise the converted
      # onnx model will have size issues.
      if (
          dynamic_axis is None
          and device.type is DeviceType.GPU
          and is_diffusion
      ):
          dynamic_axis = get_default_dynamic_info(input_sizes)
  
      if dynamic_axis is not None:
          dynamic_axis["inputs"] = [
              {int(k): v for (k, v) in val.items()}
              for val in dynamic_axis["inputs"]
          ]
          dynamic_axis["outputs"] = [
              {int(k): v for (k, v) in val.items()}
              for val in dynamic_axis["outputs"]
          ]
  
      dynamic_axis = ifnone(
          dynamic_axis,
          _extract_dynamic_axis(model, dataloader, input_sizes, device),
      )
      return batch_size, input_sizes, input_types, dynamic_axis
  
file_path: nebuly/optimization/nebullvm/nebullvm/tools/pytorch.py
imports: from pathlib import Path
  from typing import List, Tuple, Optional, Dict, Union, Sequence
  from loguru import logger
  from nebullvm.core.models import Device, DataType, DeviceType, InputInfo
  from nebullvm.optional_modules.torch import torch, DataLoader
  from nebullvm.tools.data import DataManager
  from nebullvm.tools.diffusers import get_default_dynamic_info
label_desc: Use image node_id 7 for calling a global function with example usage: extract_info_from_torch_data(model, dataloader, dynamic_axis, device, is_diffusion) and returns: batch_size,  input_sizes,  input_types,  dynamic_axis