code: def _train_step(
      model,
      optim,
      expected_param_shapes,
      amp_context,
      mixed_precision,
      half_input,
  ):
      # Prepare for training step.
      optim.zero_grad()
      model.train()
  
      # Create input and run forward pass.
      input = torch.randn(2, 3).cuda()
  
      # Make it FP16 when it is OK to do so.
      if (amp_context and half_input) or (
          mixed_precision and half_input
      ):
          input = input.half()
  
      context = contextlib.suppress()
      if amp_context:
          context = torch.cuda.amp.autocast(True)
  
      with context:
          loss = model(input).sum()
  
      _check_params(model, expected_param_shapes)
  
      # Run backward pass.
      loss.backward()
      _check_params(model, expected_param_shapes)
  
      # Finally, take a step.
      optim.step()
      _check_params(model, expected_param_shapes)
  
file_path: fairscale/tests/nn/data_parallel/test_fsdp_with_checkpoint_wrapper.py
imports: import contextlib
  import pytest
  import torch
  from torch import nn
  import torch.distributed
  import torch.multiprocessing
  from fairscale.fair_dev.testing.testing import dist_init, skip_if_single_gpu, teardown, temp_files_ctx
  from fairscale.nn.checkpoint.checkpoint_activations import checkpoint_wrapper
  from fairscale.nn.data_parallel import FullyShardedDataParallel
label_desc: Use image node_id 3 for calling a global function with example usage: _train_step(model, optim, expected_param_shapes, amp_context, mixed_precision, half_input) without return types