code: def torch_load(fn: str) -> Dict[str, Tensor]:
      t = Tensor.empty(
          os.stat(fn).st_size, dtype=dtypes.uint8, device=f"disk:{fn}"
      )
  
      offsets: Dict[Union[str, int], int] = {}
      lens: Dict[Union[str, int], int] = {}
  
      def _rebuild_tensor_v2(
          storage,
          storage_offset,
          size,
          stride,
          requires_grad=None,
          backward_hooks=None,
          metadata=None,
      ):
          # print(storage, storage_offset, size, stride, requires_grad, backward_hooks, metadata)
          lens[storage[2]] = storage[4] * storage[1].itemsize
          if storage[2] not in offsets:
              return None
          byte_offset = (
              offsets[storage[2]] + storage_offset * storage[1].itemsize
          )
          ret = t[byte_offset : byte_offset + prod(size)].cast(
              storage[1]
          )
  
          # 7 lines to deal with permuted tensors. NOTE: this currently requires reading off the disk
          shape_strides = [
              (s, st) for s, st in zip(size, stride) if s != 1
          ]
          permute_indexes = [
              len(shape_strides) - 1 - y
              for y in argsort([x[1] for x in shape_strides])
          ]
          if tuple(permute_indexes) != tuple(
              range(len(permute_indexes))
          ):
              intermediate_shape = tuple(
                  [
                      shape_strides[x][0]
                      for x in argsort(permute_indexes)
                  ]
              )
              assert tuple(
                  [
                      shape_strides[i][1]
                      for i in argsort(permute_indexes)
                  ]
              ) == strides_for_shape(
                  intermediate_shape
              ), "nonpermutable strides"
              if DEBUG >= 3:
                  print(
                      f"WARNING: this torch load is slow. CPU to permute {intermediate_shape} with {permute_indexes}"
                  )
              assert (
                  storage[1] != dtypes.bfloat16
              ), "can't CPU permute BF16"
              # TODO: find a nice way to support all shapetracker on disktensors
              ret = (
                  ret.cpu()
                  .reshape(intermediate_shape)
                  .permute(permute_indexes)
              )
  
          return ret.reshape(size)
  
      class Parameter:
          def __setstate__(self, state):
              self.tensor = state[0]
  
      deserialized_objects: Dict[str, Any] = {}
      intercept = {
          "HalfStorage": dtypes.float16,
          "FloatStorage": dtypes.float32,
          "BFloat16Storage": dtypes.bfloat16,
          "IntStorage": dtypes.int32,
          "LongStorage": dtypes.int64,
          "_rebuild_tensor_v2": _rebuild_tensor_v2,
          "FloatTensor": None,
          "Parameter": Parameter,
      }
      whitelist = {
          "torch",
          "collections",
          "numpy",
          "_codecs",
      }  # NOTE: this is not for security, only speed
  
      class Dummy:
          pass
  
      class TorchPickle(pickle.Unpickler):
          def find_class(self, module, name):
              module_root = module.split(".")[0]
              if module_root not in whitelist:
                  if DEBUG >= 2:
                      print(
                          f"WARNING: returning Dummy for {module} {name}"
                      )
                  return Dummy
              return (
                  intercept[name]
                  if module_root == "torch"
                  else super().find_class(module, name)
              )
  
          def persistent_load(self, pid):
              return deserialized_objects.get(pid, pid)
  
      if zipfile.is_zipfile(fn):
          myzip = zipfile.ZipFile(fn, "r")
          base_name = myzip.namelist()[0].split("/", 1)[0]
          for n in myzip.namelist():
              if n.startswith(f"{base_name}/data/"):
                  with myzip.open(n) as myfile:
                      offsets[n.split("/")[-1]] = myfile._orig_compress_start  # type: ignore
          with myzip.open(f"{base_name}/data.pkl") as myfile:
              return TorchPickle(myfile).load()
      elif tarfile.is_tarfile(fn):
          with tarfile.open(fn, "r") as tar:
              storages_offset = tar.getmember("storages").offset_data
              f = unwrap(tar.extractfile("storages"))
              for i in range(TorchPickle(f).load()):  # num_storages
                  (key, _, storage_type), sz = (
                      TorchPickle(f).load(),
                      struct.unpack("<q", f.read(8))[0],
                  )
                  offsets[key] = storages_offset + f.tell()
                  f.seek(sz * storage_type.itemsize, 1)
              f = unwrap(tar.extractfile("tensors"))
              for _ in range(TorchPickle(f).load()):  # num_tensors
                  (key, storage_id, _), ndim, _ = (
                      TorchPickle(f).load(),
                      struct.unpack("<i", f.read(4))[0],
                      f.read(4),
                  )
                  size, stride = struct.unpack(
                      f"<{ndim}q", f.read(8 * ndim)
                  ), struct.unpack(f"<{ndim}q", f.read(8 * ndim))
                  storage_offset = struct.unpack("<q", f.read(8))[0]
                  deserialized_objects[str(key)] = _rebuild_tensor_v2(
                      (None, storage_type, storage_id, None, -1),
                      storage_offset,
                      size,
                      stride,
                  )
              return {
                  k: v.tensor if isinstance(v, Parameter) else v
                  for k, v in TorchPickle(
                      unwrap(tar.extractfile("pickle"))
                  )
                  .load()
                  .items()
              }
      else:
          with open(fn, "rb") as f:
              pkl = TorchPickle(f)
              _, _, _, rwd, _, ids, base_offset = (
                  pkl.load(),
                  pkl.load(),
                  pkl.load(),
                  f.tell(),
                  pkl.load(),
                  pkl.load(),
                  f.tell(),
              )
              for i in ids:
                  offsets[i] = base_offset + 8
                  base_offset += 8 + lens[i]
              f.seek(rwd)
              return TorchPickle(f).load()
  
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 7 for calling a global function with example usage: torch_load(fn) and returns: ret, None, unknown, deserialized_objects, TorchPickle, Dummy, unknown, TorchPickle