code: def export(
      variant: str, output_dir: str, checkpoint: str = None
  ) -> None:
      """Method exports coreml package for mobile inference.
  
      Args:
          variant: FastViT model variant.
          output_dir: Path to save exported model.
          checkpoint: Path to trained checkpoint. Default: ``None``
      """
      # Create output directory.
      os.makedirs(output_dir, exist_ok=True)
  
      # Random input tensor for tracing purposes.
      inputs = torch.rand(1, 3, 256, 256)
      inputs_tensor = [
          coremltools.TensorType(
              name="images",
              shape=inputs.shape,
          )
      ]
  
      # Instantiate model variant.
      model = getattr(models, variant)()
      print(f"Export and Convert Model: {variant}")
  
      # Always reparameterize before exporting.
      reparameterized_model = reparameterize_model(model)
      if checkpoint is not None:
          print(f"Load checkpoint {checkpoint}")
          chkpt = torch.load(checkpoint)
          reparameterized_model.load_state_dict(chkpt["state_dict"])
      reparameterized_model.eval()
  
      # Trace and export.
      traced_model = torch.jit.trace(
          reparameterized_model, torch.Tensor(inputs)
      )
      output_path = os.path.join(output_dir, variant)
      pt_name = output_path + ".pt"
      traced_model.save(pt_name)
      ml_model = coremltools.convert(
          model=pt_name,
          outputs=None,
          inputs=inputs_tensor,
          convert_to="mlprogram",
          debug=False,
      )
      ml_model.save(output_path + ".mlpackage")
  
imports: import argparse
  import os
  import coremltools
  import torch
  import models
  from models.modules.mobileone import reparameterize_model
label_desc: Use image node_id 2 for calling a global function with example usage: export(variant, output_dir, checkpoint) without return types