code: def train(USE_PRETRAINED_MODEL=False):
      net = TitanicSimpleNNModel()
      (
          train_features,
          train_labels,
          test_features,
          test_labels,
          feature_names,
      ) = prepare()
      USE_PRETRAINED_MODEL = dict_args["use_pretrained_model"]
      if USE_PRETRAINED_MODEL:
          net.load_state_dict(
              torch.load("models/titanic_state_dict.pt")
          )
          net.eval()
          print("Model Loaded!")
      else:
          criterion = nn.CrossEntropyLoss()
          num_epochs = dict_args["max_epochs"]
          mlflow.log_param("epochs", num_epochs)
          mlflow.log_param("lr", dict_args["lr"])
  
          optimizer = torch.optim.Adam(
              net.parameters(), lr=dict_args["lr"]
          )
          print(train_features.dtype)
          input_tensor = torch.from_numpy(train_features).type(
              torch.FloatTensor
          )
          label_tensor = torch.from_numpy(train_labels)
          for epoch in range(num_epochs):
              output = net(input_tensor)
              loss = criterion(output, label_tensor)
              optimizer.zero_grad()
              loss.backward()
              optimizer.step()
              if epoch % 50 == 0:
                  print(
                      f"Epoch {epoch + 1}/{num_epochs} => Train Loss: {loss.item():.2f}"
                  )
                  mlflow.log_metric(
                      f"Epoch {epoch + 1!s} Loss",
                      float(loss.item()),
                      step=epoch,
                  )
          if not os.path.isdir("models"):
              os.makedirs("models")
              torch.save(
                  net.state_dict(), "models/titanic_state_dict.pt"
              )
      summary, _ = count_model_parameters(net)
      mlflow.log_text(str(summary), "model_summary.txt")
      return (
          net,
          train_features,
          train_labels,
          test_features,
          test_labels,
          feature_names,
      )
  
imports: import os
  from argparse import ArgumentParser
  import matplotlib.pyplot
  import numpy
  import pandas
  import torch
  from captum.attr import IntegratedGradients, LayerConductance, NeuronConductance
  from prettytable import PrettyTable
  from scipy import stats
  from sklearn.model_selection import train_test_split
  from torch import nn
  import mlflow
label_desc: Use image node_id 5 for calling a global function with example usage: train(USE_PRETRAINED_MODEL) and returns: net,  train_features,  train_labels,  test_features,  test_labels,  feature_names