code: def prune_and_save_model(model, amount):
      for _, module in model.named_modules():
          if isinstance(module, (torch.nn.Conv2d, torch.nn.Linear)):
              prune.l1_unstructured(
                  module, name="weight", amount=amount
              )
              prune.remove(module, "weight")
  
      mlflow.pytorch.save_state_dict(model.state_dict(), ".")
      model = torch.load("state_dict.pth")
      os.remove("state_dict.pth")
      return model
  
imports: import argparse
  import copy
  import os
  import shutil
  import tempfile
  from pathlib import Path
  import lightning
  import torch
  from ax.service.ax_client import AxClient
  from mnist import LightningMNISTClassifier, MNISTDataModule
  from prettytable import PrettyTable
  from torch.nn.utils import prune
  import mlflow.pytorch
  from mlflow.tracking.artifact_utils import _download_artifact_from_uri
label_desc: Use image node_id 5 for calling the IterativePrune obj's underlying member method code with example usage: obj.prune_and_save_model(model, amount) and returns: model