code: def load_state_dict(
      model,
      state_dict: Dict[str, Tensor],
      strict=True,
      verbose=True,
      consume=False,
  ):
      start_mem_used = GlobalCounters.mem_used
      with Timing(
          "loaded weights in ",
          lambda et_ns: f", {(GlobalCounters.mem_used-start_mem_used)/1e9:.2f} GB loaded at {(GlobalCounters.mem_used-start_mem_used)/et_ns:.2f} GB/s",
      ):  # noqa: E501
          model_state_dict = get_state_dict(model)
          if DEBUG >= 1 and len(state_dict) > len(model_state_dict):
              print(
                  "WARNING: unused weights in state_dict",
                  sorted(
                      list(state_dict.keys() - model_state_dict.keys())
                  ),
              )
          for k, v in (
              t := tqdm(
                  model_state_dict.items(), disable=CI or not verbose
              )
          ):
              t.set_description(
                  f"ram used: {GlobalCounters.mem_used/1e9:5.2f} GB, {k:50s}"
              )
              if k not in state_dict and not strict:
                  if DEBUG >= 1:
                      print(f"WARNING: not loading {k}")
                  continue
              v.assign(
                  state_dict[k].shard(mlb.device, mlb.axis)
                  if isinstance((mlb := v.lazydata), MultiLazyBuffer)
                  else state_dict[k].to(v.device)
              ).realize()
              if consume:
                  del state_dict[k]
  
file_path: tinygrad/tinygrad/nn/state.py
imports: import os, json, pathlib, zipfile, pickle, tarfile, struct
  from tqdm import tqdm
  from typing import Dict, Union, List, Optional, Any, Tuple
  from tinygrad.tensor import Tensor
  from tinygrad.ops import GlobalCounters
  from tinygrad.dtype import dtypes
  from tinygrad.helpers import prod, argsort, DEBUG, Timing, CI, unwrap
  from tinygrad.shape.view import strides_for_shape
  from tinygrad.features.multi import MultiLazyBuffer
  from collections import OrderedDict
label_desc: Use image node_id 6 for calling a global function with example usage: load_state_dict(model, state_dict, strict, verbose, consume) without return types