code: def setup_data(
      train_dataset_path: components.OutputPath("tensor"),  # noqa: F821
      test_dataset_path: components.OutputPath("tensor"),  # noqa: F821
  ):
      import torch
      from torchvision import datasets, transforms
  
      transform = transforms.Compose(
          [
              transforms.ToTensor(),
              transforms.Normalize((0.1307,), (0.3081,)),
          ]
      )
      train_dataset = datasets.FakeData(
          size=2000,
          image_size=(1, 28, 28),
          num_classes=10,
          transform=transform,
      )
      test_dataset = datasets.FakeData(
          size=2000,
          image_size=(1, 28, 28),
          num_classes=10,
          transform=transform,
      )
  
      torch.save(train_dataset, train_dataset_path)
      torch.save(test_dataset, test_dataset_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 2 for calling a global function with example usage: setup_data(train_dataset_path, test_dataset_path) without return types