code: def run_test_train_eval_change(rank, world_size, file):
      # Check that ShardedDDP handles the switch from training to eval properly
      dist.init_process_group(
          init_method="file://" + file,
          backend="gloo",
          rank=rank,
          world_size=world_size,
      )
  
      model = _get_mlp()
      model.train()
      optimizer = OSS(
          params=model.parameters(),
          optim=torch.optim.SGD,
          lr=1e-3,
          momentum=0.99,
      )
      model = ShardedDataParallel(model, optimizer)
      input_tensor = torch.rand((2, 2))
      loss = model(input_tensor).sum()
      loss.backward()  # make sure that the gradients are reduced
  
      # Wipe the gradients and switch to eval mode
      model.zero_grad()
      model.eval()
      _ = model(input_tensor)
      assert (
          next(model.parameters()).grad is None
          or torch.norm(next(model.parameters()).grad) < 1e-6
      )
  
      # Get back to training
      model = model.train()
      model(input_tensor).sum().backward()
      assert torch.norm(next(model.parameters()).grad) > 0.0
  
      dist.destroy_process_group()
  
imports: from contextlib import suppress
  import numpy
  import pytest
  import torch
  import torch.distributed
  import torch.multiprocessing
  from torch.nn import Linear, Sequential
  from fairscale.fair_dev.testing.testing import GPT2, SGDWithPausingCompute, available_devices, check_same_models_across_ranks, skip_if_less_than_four_gpu, skip_if_no_cuda, skip_if_single_gpu, temp_files_ctx
  from fairscale.nn.data_parallel import ShardedDataParallel
  from fairscale.optim import OSS
label_desc: Use image node_id 11 for calling a global function with example usage: run_test_train_eval_change(rank, world_size, file) without return types