code: def test_full_pytorch_example(large_mock_mnist_data, tmpdir):
      # First, generate mock dataset
      dataset_url = "file://{}".format(tmpdir)
      mnist_data_to_petastorm_dataset(
          tmpdir,
          dataset_url,
          mnist_data=large_mock_mnist_data,
          spark_master="local[1]",
          parquet_files_count=1,
      )
  
      # Next, run a round of training using the pytorce adapting data loader
      from petastorm.pytorch import DataLoader
  
      torch.manual_seed(1)
      device = torch.device("cpu")
      model = pytorch_example.Net().to(device)
      optimizer = torch.optim.SGD(
          model.parameters(), lr=0.01, momentum=0.5
      )
  
      transform = TransformSpec(
          pytorch_example._transform_row, removed_fields=["idx"]
      )
  
      with DataLoader(
          make_reader(
              "{}/train".format(dataset_url),
              reader_pool_type="dummy",
              num_epochs=1,
              transform_spec=transform,
          ),
          batch_size=32,
      ) as train_loader:
          pytorch_example.train(
              model, device, train_loader, 10, optimizer, 1
          )
      with DataLoader(
          make_reader(
              "{}/test".format(dataset_url),
              reader_pool_type="dummy",
              num_epochs=1,
              transform_spec=transform,
          ),
          batch_size=100,
      ) as test_loader:
          pytorch_example.test(model, device, test_loader)
  
file_path: petastorm/examples/mnist/tests/test_pytorch_mnist.py
imports: import logging
  import os
  import unittest
  import pyarrow
  import torch
  import pytest
  import examples.mnist.pytorch_example
  from examples.mnist.generate_petastorm_mnist import mnist_data_to_petastorm_dataset, download_mnist_data
  from examples.mnist.tests.conftest import SMALL_MOCK_IMAGE_COUNT
  from petastorm import make_reader, TransformSpec
label_desc: Use image node_id 2 for calling a global function with example usage: test_full_pytorch_example(large_mock_mnist_data, tmpdir) without return types