code: def train_model(model, epochs, X_train, y_train):
      criterion = nn.CrossEntropyLoss()
      optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
  
      for epoch in range(epochs):
          out = model(X_train)
          loss = criterion(out, y_train).to(device)
          optimizer.zero_grad()
          loss.backward()
          optimizer.step()
  
          if epoch % 10 == 0:
              print("number of epoch", epoch, "loss", float(loss))
  
      return model
  
file_path: mlflow/examples/pytorch/torchscript/IrisClassification/iris_classification.py
imports: import argparse
  import torch
  import torch.nn.functional
  from sklearn.datasets import load_iris
  from sklearn.metrics import accuracy_score
  from sklearn.model_selection import train_test_split
  from torch import nn
  import mlflow.pytorch
  from mlflow.models import infer_signature
label_desc: Use image node_id 2 for calling a global function with example usage: train_model(model, epochs, X_train, y_train) and returns: model