code: def train_torch(
      move_to_device: Callable,
      precision_context,
      input_dtype=torch.float32,
  ):
      make_deterministic(warn_only=True)
      memory_stats = {}
  
      model = ConvNet()
      model = move_to_device(model)
      dataloader = model.get_dataloader()
      optimizer = model.get_optimizer()
      loss_fn = model.get_loss_function()
  
      memory_stats["start"] = torch.cuda.memory_stats()
  
      model.train()
      iteration_timings = []
      iterator = iter(dataloader)
      for _ in range(model.num_steps):
          t0 = time.perf_counter()
  
          inputs, labels = next(iterator)
          inputs, labels = move_to_device(inputs), move_to_device(
              labels
          )
          optimizer.zero_grad()
          with precision_context():
              outputs = model(inputs.to(input_dtype))
          loss = loss_fn(outputs.float(), labels)
          loss.backward()
          optimizer.step()
  
          t1 = time.perf_counter()
          iteration_timings.append(t1 - t0)
  
      memory_stats["end"] = torch.cuda.memory_stats()
  
      return (
          model.state_dict(),
          torch.tensor(iteration_timings),
          memory_stats,
      )
  
file_path: pytorch-lightning/tests/parity_fabric/test_parity_simple.py
imports: import time
  from copy import deepcopy
  from typing import Callable
  import pytest
  import torch
  import torch.distributed
  import torch.nn.functional
  from lightning.fabric.fabric import Fabric
  from tests_fabric.helpers.runif import RunIf
  from parity_fabric.models import ConvNet
  from parity_fabric.utils import cuda_reset, get_model_input_dtype, is_cuda_memory_close, is_state_dict_equal, is_timing_close, make_deterministic
label_desc: Use image node_id 1 for calling a global function with example usage: train_torch(move_to_device, precision_context, input_dtype) and returns: model,  torch,  memory_stats