code: def load_ckpt(self, path):
      self.state = load_ckpt_v2(
          self.state,
          path,
          self.state_shard,
          not self.config.get("eval_only", False),
      )
  
imports: import gc
  import random
  import time
  from functools import partial
  import haiku
  import jax
  import jax.numpy
  import numpy
  import optax
  from jax.experimental.maps import thread_resources
  from jax.experimental.pjit import pjit
  from mesh_transformer.checkpoint import read_ckpt, write_ckpt, write_ckpt_v2, load_ckpt_v2
  from mesh_transformer.layers import EmbeddingShard, TransformerLayerShard, RelativePositionEmbs, ProjectionShard, TransformerLayerShardV2, Projection, EmbeddingShardV2
  from mesh_transformer.util import to_f32, to_bf16, maybe_shard, head_print, global_norm
  from jax.experimental import PartitionSpec
label_desc: Use image node_id 3 for calling the CausalTransformerV2 obj's underlying member method code with example usage: obj.load_ckpt(path) without return types