code: def save_xla_ckpt(ckpt, file_or_path):
      """
      Similar to xm.save, but only try to convert "model" and "optimizer" in an MMF
      checkpoint to CPU, since they hold PyTorch tensors. Other items like lr_scheduler
      often cannot be saved with xm.save due to its errors in handling mappingproxy.
  
      Only save on the global main process (which is different from the default behavior
      of xm.save that saves a checkpoint on each node).
      """
      should_write_data = is_main()
  
      is_full_ckpt = (
          isinstance(ckpt, dict)
          and "model" in ckpt
          and "optimizer" in ckpt
      )
      if is_full_ckpt:
          ckpt["model"] = xm._maybe_convert_to_cpu(
              ckpt["model"], convert=should_write_data
          )
          ckpt["optimizer"] = xm._maybe_convert_to_cpu(
              ckpt["optimizer"], convert=should_write_data
          )
      else:
          ckpt = xm._maybe_convert_to_cpu(
              ckpt, convert=should_write_data
          )
  
      if should_write_data:
          torch.save(ckpt, file_or_path)
      xm.rendezvous("mmf.utils.checkpoint.save_xla_ckpt")
  
label_desc: Use image node_id 1 for calling a global function with example usage: save_xla_ckpt(ckpt, file_or_path) without return types