docstr: Update the fp16 chunks via their fp32 chunks. It's used by the optimizer.
docstr_tok: ["Update","the","fp16","chunks","via","their","fp32","chunks",".","It","'s","used","by","the","optimizer","."]
code: def optim_update(self) -> None:
      """Update the fp16 chunks via their fp32 chunks. It's used by the optimizer."""
      # sanity check
      assert self.paired_chunk is not None
  
      friend_chunk = self.paired_chunk
      if self.is_gathered is True:
          assert friend_chunk.is_gathered is True
          self.cuda_global_chunk.copy_(friend_chunk.cuda_global_chunk)
          self.optim_sync_flag = True
      elif friend_chunk.device_type in (
          "cuda",
          "npu",
      ) and self.device_type in ("cuda", "npu"):
          self.cuda_shard.copy_(friend_chunk.cuda_shard)
          self.optim_sync_flag = True
          self.cpu_vis_flag = False
      else:
          # optim_sync_flag is set to False
          # see shard_move function for more details
          assert friend_chunk.device_type == "cpu"
          assert self.device_type == "cpu"
          self.optim_sync_flag = False
          self.cpu_vis_flag = False
  
imports: from dataclasses import dataclass
  from enum import Enum
  from typing import Dict, List, Optional
  import torch
  import torch.distributed
  from torch.distributed import ProcessGroup
  from colossalai.utils import get_current_device
  from colossalai.utils.device import IS_NPU_AVAILABLE
label_desc: Use image node_id 22 for calling the Chunk obj's underlying member method code with example usage: obj.optim_update() without return types