code: def _export_model_onnx(
      model: torch.nn.Module,
      save_path: Path,
      model_name: str,
      input_batch: Tuple,
  ):
      if torch.cuda.is_available():
          input_batch = tuple(t.cuda() for t in input_batch)
          model.cuda()
  
      exporter = ModuleExporter(model, output_dir=save_path)
      with torch.no_grad():
          example_outputs = model(*input_batch)
      exporter.export_onnx(
          input_batch, name=model_name, example_outputs=example_outputs
      )
      onnx_path = save_path / model_name
  
      return onnx_path
  
imports: import json
  import logging
  import os.path
  from pathlib import Path
  from tempfile import TemporaryDirectory
  from typing import Tuple, List, Any, Dict
  import torch
  from sparseml.onnx.optim import ModelAnalyzer, pruning_loss_sens_magnitude
  from sparseml.pytorch.optim import ScheduledModifierManager
  from sparseml.pytorch.sparsification import EpochRangeModifier, GMPruningModifier
  from sparseml.pytorch.utils import ModuleExporter
  from sparsify.blueprints.utils import default_epochs_distribution, PruningModelEvaluator, default_pruning_settings
  from sparsify.schemas import ProjectModelAnalysisSchema
  from torch.nn import CrossEntropyLoss, MSELoss
  from torch.optim import SGD
  from tqdm.auto import tqdm
label_desc: Use image node_id 1 for calling a global function with example usage: _export_model_onnx(model, save_path, model_name, input_batch) and returns: onnx_path