MCPcopy Create free account
hub / github.com/tdrussell/diffusion-pipe / reduce_storage

Function reduce_storage

utils/reduction.py:598–626  ·  view source on GitHub ↗
(storage)

Source from the content-addressed store, hash-verified

596
597
598def reduce_storage(storage):
599 if storage.is_cuda:
600 raise RuntimeError(
601 "Cannot pickle CUDA storage; try pickling a CUDA tensor instead"
602 )
603 elif storage.device.type == "meta":
604 raise RuntimeError(
605 "Cannot pickle meta storage; try pickling a meta tensor instead"
606 )
607 elif torch.multiprocessing.get_sharing_strategy() == "file_system":
608 metadata = storage._share_filename_cpu_()
609 cache_key = metadata[1]
610 rebuild = rebuild_storage_filename
611 if isinstance(storage, torch.TypedStorage):
612 metadata += (storage.dtype,)
613 storage._shared_incref()
614 elif storage.size() == 0:
615 # This is special cased because Empty tensors
616 # (with size 0) cannot be mmapped.
617 return (rebuild_storage_empty, (type(storage),))
618 else:
619 fd, size = storage._share_fd_cpu_()
620 df = multiprocessing.reduction.DupFd(fd)
621 cache_key = fd_id(fd)
622 metadata = (df, size)
623 rebuild = rebuild_storage_fd # type: ignore[assignment]
624
625 shared_cache[cache_key] = StorageWeakRef(storage)
626 return (rebuild, (type(storage),) + metadata)
627
628
629def init_reductions():

Callers

nothing calls this directly

Calls 2

fd_idFunction · 0.85
StorageWeakRefClass · 0.85

Tested by

no test coverage detected