code: def _worker(
      gpu_id: int, sync_file: str, world_size: int, with_heads: bool
  ):
      torch.manual_seed(0)
      os.environ["RANK"] = str(gpu_id)
      init_distributed_on_file(
          world_size=world_size, gpu_id=gpu_id, sync_file=sync_file
      )
      torch.backends.cudnn.deterministic = True
  
      config = TestCheckpointConversion._create_fsdp_model_config(
          with_fsdp=True
      )
      model = build_model(config.MODEL, config.OPTIMIZER).cuda(gpu_id)
      model = fsdp_wrapper(model, **config.MODEL.FSDP_CONFIG)
      optimizer = optim.SGD(model.parameters(), lr=1e-4)
  
      # Fake inputs
      num_iterations = 5
      batch_size = 3
      torch.manual_seed(gpu_id)
      fake_inputs = torch.randn(
          size=(num_iterations, batch_size, 3, 96, 96)
      )
      fake_targets = torch.randn(size=(num_iterations, batch_size))
  
      # Fake training loop
      criterion = nn.MSELoss()
      for iteration in range(num_iterations):
          fake_input = fake_inputs[iteration].cuda(gpu_id)
          fake_target = fake_targets[iteration].cuda(gpu_id)
          output1, output2 = model(fake_input)[0]
          loss = criterion(
              output1.sum(axis=-1), fake_target
          ) + criterion(output2.sum(axis=-1), fake_target)
          if gpu_id == 0:
              print(loss.item())
          optimizer.zero_grad()
          loss.backward()
          optimizer.step()
  
      # Save a bunch of checkpoint, one by shard
      checkpoint_writer = CheckpointWriter(
          checkpoint_folder=".",
          is_final_train_phase=True,
          mode="iteration",
          mode_num=0,
          backend="disk",
      )
      content = {
          "classy_state_dict": {
              "base_model": {
                  "model": {
                      "trunk": model.trunk.local_state_dict(),
                      "heads": [
                          head.local_state_dict()
                          for head in model.heads
                      ],
                  },
                  "meta": {
                      "trunk": model.trunk.local_metadata_dict(),
                      "heads": [
                          head.local_metadata_dict()
                          for head in model.heads
                      ],
                  },
              }
          }
      }
  
      if not with_heads:
          del content["classy_state_dict"]["base_model"]["model"][
              "heads"
          ]
          del content["classy_state_dict"]["base_model"]["meta"][
              "heads"
          ]
  
      checkpoint_writer.save_sharded_checkpoint(
          content, shard_rank=gpu_id, world_size=world_size
      )
      dist.barrier()
      print(os.listdir("."))
  
      # Convert the checkpoint to consolidated and sliced checkpoints
      if gpu_id == 0:
          CheckpointFormatConverter.sharded_to_consolidated_checkpoint(
              "checkpoint.torch", "checkpoint_conso.torch"
          )
          CheckpointFormatConverter.sharded_to_sliced_checkpoint(
              "checkpoint.torch", "checkpoint_sliced.torch"
          )
      dist.barrier()
      print(os.listdir("."))
  
      # Now create models initialized from the previous checkpoint and compare them
      fake_test_input = torch.randn(size=(1, 3, 96, 96)).cuda(gpu_id)
  
      shard_cp = CheckpointLoader.load_and_broadcast_init_weights(
          "checkpoint.torch", device=torch.device("cpu")
      )
      shard_model = build_model(config.MODEL, config.OPTIMIZER).cuda(
          gpu_id
      )
      shard_model = fsdp_wrapper(
          shard_model, **config.MODEL.FSDP_CONFIG
      )
      shard_model.init_model_from_weights_params_file(config, shard_cp)
  
      conso_cp = CheckpointLoader.load_and_broadcast_init_weights(
          "checkpoint_conso.torch", device=torch.device("cpu")
      )
      conso_model = build_model(config.MODEL, config.OPTIMIZER).cuda(
          gpu_id
      )
      conso_model = fsdp_wrapper(
          conso_model, **config.MODEL.FSDP_CONFIG
      )
      conso_model.init_model_from_weights_params_file(config, conso_cp)
  
      slice_cp = CheckpointLoader.load_and_broadcast_init_weights(
          "checkpoint_sliced.torch", device=torch.device("cpu")
      )
      slice_model = build_model(config.MODEL, config.OPTIMIZER).cuda(
          gpu_id
      )
      slice_model = fsdp_wrapper(
          slice_model, **config.MODEL.FSDP_CONFIG
      )
      slice_model.init_model_from_weights_params_file(config, slice_cp)
  
      # Verifying that the models are equivalent
      if gpu_id == 0:
          slice_state_dict = slice_model.local_state_dict()
          conso_state_dict = conso_model.local_state_dict()
          assert set(slice_state_dict.keys()) == set(
              conso_state_dict.keys()
          )
          for k in slice_state_dict.keys():
              slice_val = slice_state_dict[k]
              conso_val = conso_state_dict[k]
              assert torch.allclose(
                  slice_val, conso_val
              ), f"Difference for key {k}: {slice_val} VS {conso_val}"
      dist.barrier()
  
      with torch.no_grad():
          ref_out = model.trunk(fake_test_input)[0]
          shard_out = shard_model.trunk(fake_test_input)[0]
          conso_out = conso_model.trunk(fake_test_input)[0]
          slice_out = slice_model.trunk(fake_test_input)[0]
          assert torch.allclose(
              ref_out, shard_out
          ), f"{ref_out.sum()} vs {shard_out.sum()}"
          assert torch.allclose(
              ref_out, conso_out
          ), f"{ref_out.sum()} vs {conso_out.sum()}"
          assert torch.allclose(
              ref_out, slice_out
          ), f"{ref_out.sum()} vs {slice_out.sum()}"
  
file_path: vissl/tests/test_state_checkpoint_conversion.py
imports: import os
  import unittest
  import torch
  import torch.distributed
  import torch.multiprocessing
  import torch.nn
  import torch.optim
  from vissl.models import build_model
  from vissl.utils.checkpoint import CheckpointFormatConverter, CheckpointLoader, CheckpointWriter
  from vissl.utils.fsdp_utils import fsdp_wrapper
  from vissl.utils.hydra_config import compose_hydra_configuration, convert_to_attrdict
  from vissl.utils.test_utils import gpu_test, in_temporary_directory, init_distributed_on_file, with_temp_files
label_desc: Use image node_id 2 for calling the TestCheckpointConversion obj's underlying member method code with example usage: obj._worker(gpu_id, sync_file, world_size, with_heads) without return types