Method create_tiny_mapping

extra/torch_hook/hook_torch.py:44–47

(arg)

      return args is None or isinstance(arg, (str, int, float, bool))

    def create_tiny_mapping(arg):
      if WRAP_TINY:
        tt = Tensor.from_blob(arg.data_ptr(), arg.shape, dtype=_from_torch_dtype(arg.dtype))
        cuda_to_tiny_mappings[arg] = tiny_torch.wrap(tt)

    for i,arg in enumerate(args):
      if torch.is_tensor(arg):