code: def train_model(
      train_dataloader_path: components.InputPath(
          "dataloader"
      ),  # noqa: F821
      test_dataloader_path: components.InputPath(
          "dataloader"
      ),  # noqa: F821
      model_path: components.OutputPath("pytorch_model"),  # noqa: F821
      seed: int = 1337,
      use_cuda: bool = False,
      lr: float = 1.0,
      gamma: float = 0.7,
      epochs: int = 1,
      log_interval: int = 10,
      dry_run: bool = False,
      save_model: bool = False,
  ):
      import torch
      import torch.nn as nn
      import torch.nn.functional as F  # noqa: N812
      import torch.optim as optim
      from torch.optim.lr_scheduler import StepLR
  
      class Net(nn.Module):
          def __init__(self):
              super().__init__()
              self.conv1 = nn.Conv2d(1, 32, 3, 1)
              self.conv2 = nn.Conv2d(32, 64, 3, 1)
              self.dropout1 = nn.Dropout(0.25)
              self.dropout2 = nn.Dropout(0.5)
              self.fc1 = nn.Linear(9216, 128)
              self.fc2 = nn.Linear(128, 10)
  
          def forward(self, x):
              x = self.conv1(x)
              x = F.relu(x)
              x = self.conv2(x)
              x = F.relu(x)
              x = F.max_pool2d(x, 2)
              x = self.dropout1(x)
              x = torch.flatten(x, 1)
              x = self.fc1(x)
              x = F.relu(x)
              x = self.dropout2(x)
              x = self.fc2(x)
              output = F.log_softmax(x, dim=1)
              return output
  
      def train(
          model,
          device,
          train_loader,
          optimizer,
          epoch,
          log_interval,
          dry_run,
      ):
          model.train()
          for batch_idx, (data, target) in enumerate(train_loader):
              data, target = data.to(device), target.to(device)
              optimizer.zero_grad()
              output = model(data)
              loss = F.nll_loss(output, target)
              loss.backward()
              optimizer.step()
              if batch_idx % log_interval == 0:
                  wandb.log(
                      {
                          "epoch": epoch,
                          "step": batch_idx * len(data),
                          "loss": loss.item(),
                      }
                  )
                  if dry_run:
                      break
  
      def test(model, device, test_loader):
          model.eval()
          test_loss = 0
          correct = 0
          with torch.no_grad():
              for data, target in test_loader:
                  data, target = data.to(device), target.to(device)
                  output = model(data)
                  test_loss += F.nll_loss(
                      output, target, reduction="sum"
                  ).item()  # sum up batch loss
                  pred = output.argmax(
                      dim=1, keepdim=True
                  )  # get the index of the max log-probability
                  correct += pred.eq(target.view_as(pred)).sum().item()
  
          test_loss /= len(test_loader.dataset)
          wandb.log(
              {
                  "test_loss": test_loss,
                  "accuracy": correct / len(test_loader.dataset),
              }
          )
  
      torch.manual_seed(seed)
      train_loader = torch.load(train_dataloader_path)
      test_loader = torch.load(test_dataloader_path)
  
      device = torch.device("cuda" if use_cuda else "cpu")
      model = Net()
      model.to(device)
      optimizer = optim.Adadelta(model.parameters(), lr=lr)
      scheduler = StepLR(optimizer, step_size=1, gamma=gamma)
      for epoch in range(1, epochs + 1):
          train(
              model,
              device,
              train_loader,
              optimizer,
              epoch,
              log_interval,
              dry_run,
          )
          test(model, device, test_loader)
          scheduler.step()
  
      if save_model:
          torch.save(model.state_dict(), model_path)
  
file_path: wandb/tests/functional_tests/t0_main/kfp/kfp-pipeline-pytorch.py
imports: import os
  import random
  import kfp
  import kfp.dsl
  import wandb
  from kfp import components
  from kubernetes.client.models import V1EnvVar
  from wandb.integration.kfp import wandb_log
  from wandb_probe import wandb_probe_package
label_desc: Use image node_id 4 for calling a global function with example usage: train_model(train_dataloader_path, test_dataloader_path, model_path, seed, use_cuda, lr, gamma, epochs, log_interval, dry_run, save_model) and returns: output