From 77411ca653aba0acde18dda46a1ff3d09d44f25c Mon Sep 17 00:00:00 2001 From: Kyle Sayers Date: Mon, 16 Dec 2024 14:12:34 -0500 Subject: [PATCH] add docstring Signed-off-by: Kyle Sayers --- src/compressed_tensors/utils/offload.py | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/src/compressed_tensors/utils/offload.py b/src/compressed_tensors/utils/offload.py index f6ef4830..7110a250 100644 --- a/src/compressed_tensors/utils/offload.py +++ b/src/compressed_tensors/utils/offload.py @@ -263,6 +263,17 @@ def offload_to_weights_map( value: torch.Tensor, default_device: torch.device = torch.device("cpu"), ): + """ + Helper function which implements offloaded item assignment for PrefixedDataset, + OffloadedWeightsLoader, and Dict types. + + :param weights_map: weight map to be updated with offload information + :param key: key used to identify weight location + :param value: weight being offloaded + :param default_device: in the event that the weights_map does already contain + offloaded weights or use disk offloading, the weight will be offloaded to the + `default_device` + """ if isinstance(weights_map, PrefixedDataset): dataset = weights_map.dataset key = f"{weights_map.prefix}{key}"