code: def prepare_optimizer(model: torch.nn.Module):
      optimize_params = [
          param
          for param in model.parameters()
          if param.requires_grad == True
      ]
      optimizer = nni.trace(Adam)(optimize_params, lr=0.001)
      return optimizer
  
file_path: nni/examples/compression/models.py
imports: from pathlib import Path
  import torch
  import torch.nn.functional
  from torch.optim import Adam
  from torch.optim.lr_scheduler import _LRScheduler
  from torch.utils.data import DataLoader
  from torchvision import datasets, transforms
  from torchvision.models.mobilenetv3 import mobilenet_v3_small
  from torchvision.models.resnet import resnet18
  import nni
label_desc: Use image node_id 4 for calling a global function with example usage: prepare_optimizer(model) and returns: optimizer