code: def shard_on_the_fly(
      switch_checkpoint_path,
      dump_path,
      num_experts,
      dtype,
      weights_name: str = WEIGHTS_NAME,
  ):
      sharded_state_dicts = []
      total_size = 0
      os.makedirs(dump_path, exist_ok=True)
  
      for expert in range(num_experts):
          expert_path = switch_checkpoint_path + f"-rank-{expert}.pt"
          if os.path.isfile(expert_path):
              expert_state = torch.load(expert_path)["model"]
              remove_ignore_keys_(expert_state)
              expert_state = rename_fairseq_keys(expert_state, expert)
              save_path = os.path.join(
                  dump_path,
                  weights_name.replace(
                      ".bin",
                      f"-{len(sharded_state_dicts)+1:05d}-of-???.bin",
                  ),
              )
              torch.save(expert_state, save_path)
              sharded_state_dicts.append(expert_state.keys())
              total_size += sum(
                  [value.numel() for key, value in expert_state.items()]
              ) * dtype_byte_size(
                  expert_state[list(expert_state)[0]].dtype
              )
  
      # Add the last block
      save_path = os.path.join(
          dump_path,
          weights_name.replace(
              ".bin", f"-{len(sharded_state_dicts)+1:05d}-of-???.bin"
          ),
      )
      shared_weights = torch.load(
          switch_checkpoint_path + "-shared.pt"
      )["model"]
      remove_ignore_keys_(shared_weights)
      shared_weights = rename_fairseq_keys(shared_weights, None)
      shared_weights["shared.weight"] = shared_weights[
          "decoder.embed_tokens.weight"
      ]
      sharded_state_dicts.append(shared_weights.keys())
  
      # If we only have the shared weights (dummy model/experts saved on the same file)
      if len(sharded_state_dicts) == 1:
          save_path = os.path.join(dump_path, weights_name)
          torch.save(shared_weights, save_path)
          return {weights_name: sharded_state_dicts[0]}, None
      else:
          torch.save(shared_weights, save_path)
      # Otherwise, let's build the index
      weight_map = {}
      for idx, shard in enumerate(sharded_state_dicts):
          shard_file = weights_name.replace(
              ".bin",
              f"-{idx+1:05d}-of-{len(sharded_state_dicts):05d}.bin",
          )
          temp_filename = os.path.join(
              dump_path,
              weights_name.replace(".bin", f"-{idx+1:05d}-of-???.bin"),
          )
          os.rename(temp_filename, os.path.join(dump_path, shard_file))
          for key in shard:
              weight_map[key] = shard_file
  
      # Add the metadata
      metadata = {"total_size": total_size}
      index = {"metadata": metadata, "weight_map": weight_map}
  
      with open(
          os.path.join(dump_path, WEIGHTS_INDEX_NAME),
          "w",
          encoding="utf-8",
      ) as f:
          content = json.dumps(index, indent=2, sort_keys=True) + "\n"
          f.write(content)
  
      return metadata, index
  
file_path: transformers/src/transformers/models/nllb_moe/convert_nllb_moe_sharded_original_checkpoint_to_pytorch.py
imports: import argparse
  import json
  import os
  import torch
  from torch import nn
  from transformers import NllbMoeConfig, NllbMoeModel
  from transformers.modeling_utils import dtype_byte_size
  from transformers.utils import WEIGHTS_INDEX_NAME, WEIGHTS_NAME
label_desc: Use image node_id 4 for calling a global function with example usage: shard_on_the_fly(switch_checkpoint_path, dump_path, num_experts, dtype, weights_name) and returns: metadata,  index, dict,  None