compressed-tensors 0.9.4a20250421__py3-none-any.whl → 0.9.5a20250424__py3-none-any.whl

compressed_tensors/utils/offload.py

@@ -94,22 +94,6 @@ def is_module_offloaded(module: torch.nn.Module) -> bool:
     return has_offloaded_params(module)
 
 
-def get_execution_device(module: torch.nn.Module) -> torch.device:
-    """
-    :param module: module to check
-    :return: device module is loaded onto during forward pass
-    """
-    if has_offloaded_params(module):
-        return module._hf_hook.execution_device
-    device = next(module.parameters()).device
-
-    # offload only gets set for leaf modules, fallback to checking for device type
-    if device.type == "meta":
-        return module._hf_hook.execution_device
-
-    return device
-
-
 def get_offloaded_device(module: torch.nn.Module) -> torch.device:
     """
     :param module: module to check
@@ -158,6 +142,26 @@ def update_parameter_data(
 """ Candidates for Upstreaming """
 
 
+def get_execution_device(module: torch.nn.Module) -> torch.device:
+    """
+    Get the device which inputs should be moved to before module execution
+
+    :param module: module to check, may be offloaded
+    :return: onload device of module
+    """
+    if has_offloaded_params(module):
+        return module._hf_hook.execution_device
+
+    first_param = next(module.parameters(), None)
+    if first_param is None:
+        warnings.warn(
+            f"Unable able to infer execution device of {module}, falling back to CPU"
+        )
+        return torch.device("cpu")
+
+    return first_param.device
+
+
 def register_offload_parameter(
     module: torch.nn.Module,
     name: str,

compressed_tensors/version.py

@@ -17,5 +17,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.9.4a20250421'
-__version_tuple__ = version_tuple = (0, 9, 4, 'a20250421')
+__version__ = version = '0.9.5.a20250424'
+__version_tuple__ = version_tuple = (0, 9, 5)

{compressed_tensors-0.9.4a20250421.dist-info → compressed_tensors-0.9.5a20250424.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: compressed-tensors
-Version: 0.9.4a20250421
+Version: 0.9.5a20250424
 Summary: Library for utilization of compressed safetensors of neural network models
 Home-page: https://github.com/neuralmagic/compressed-tensors
 Author: Neuralmagic, Inc.

{compressed_tensors-0.9.4a20250421.dist-info → compressed_tensors-0.9.5a20250424.dist-info}/RECORD

@@ -1,6 +1,6 @@
 compressed_tensors/__init__.py,sha256=UtKmifNeBCSE2TZSAfduVNNzHY-3V7bLjZ7n7RuXLOE,812
 compressed_tensors/base.py,sha256=73HYH7HY7O2roC89yG_piPFnZwrBfn_i7HmKl90SKc0,875
-compressed_tensors/version.py,sha256=mEeI8z7lrIGz844X2MRnlIQaYx5ruMqWHQ5T5U56XMc,533
+compressed_tensors/version.py,sha256=fMpLfUNedNFTmTmQeHxGZnMaXAKOKiqpI9xyx46F2gI,521
 compressed_tensors/compressors/__init__.py,sha256=smSygTSfcfuujRrAXDc6uZm4L_ccV1tWZewqVnOb4lM,825
 compressed_tensors/compressors/base.py,sha256=nvWsv4xEw1Tkxkxth6TmHplDYXfBeP22xWxOsZERyDY,7204
 compressed_tensors/compressors/helpers.py,sha256=OK6qxX9j3bHwF9JfIYSGMgBJe2PWjlTA3byXKCJaTIQ,5431
@@ -40,13 +40,13 @@ compressed_tensors/registry/__init__.py,sha256=FwLSNYqfIrb5JD_6OK_MT4_svvKTN_nEh
 compressed_tensors/registry/registry.py,sha256=vRcjVB1ITfSbfYUaGndBBmqhip_5vsS62weorVg0iXo,11896
 compressed_tensors/utils/__init__.py,sha256=gS4gSU2pwcAbsKj-6YMaqhm25udFy6ISYaWBf-myRSM,808
 compressed_tensors/utils/helpers.py,sha256=RrNvzD08naEjEiXdU-FdZjQVda1nQywu1hA_GCDj0vg,10415
-compressed_tensors/utils/offload.py,sha256=fXi_d0pcsE1X5Lx6wtKwy2CFM0G4bgeUo-zPA13XIOQ,14047
+compressed_tensors/utils/offload.py,sha256=Fmb4jBJhH5OdSQFaecFSHK_UreSyZdynEkadZ_oKcvM,14153
 compressed_tensors/utils/permutations_24.py,sha256=kx6fsfDHebx94zsSzhXGyCyuC9sVyah6BUUir_StT28,2530
 compressed_tensors/utils/permute.py,sha256=V6tJLKo3Syccj-viv4F7ZKZgJeCB-hl-dK8RKI_kBwI,2355
 compressed_tensors/utils/safetensors_load.py,sha256=rwj0ufU5561ScWDoCG7tzLBRDtiykNno2Iq4PM_JA7E,11499
 compressed_tensors/utils/semi_structured_conversions.py,sha256=XKNffPum54kPASgqKzgKvyeqWPAkair2XEQXjkp7ho8,13489
-compressed_tensors-0.9.4a20250421.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-compressed_tensors-0.9.4a20250421.dist-info/METADATA,sha256=15BwiWOfGGmgpKBoUpgBSltD_hMfMpXJBBgWLBPLzxo,7004
-compressed_tensors-0.9.4a20250421.dist-info/WHEEL,sha256=pxyMxgL8-pra_rKaQ4drOZAegBVuX-G_4nRHjjgWbmo,91
-compressed_tensors-0.9.4a20250421.dist-info/top_level.txt,sha256=w2i-GyPs2s1UwVxvutSvN_lM22SXC2hQFBmoMcPnV7Y,19
-compressed_tensors-0.9.4a20250421.dist-info/RECORD,,
+compressed_tensors-0.9.5a20250424.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+compressed_tensors-0.9.5a20250424.dist-info/METADATA,sha256=P0oAhrS28ZU90nUEi9yjIu3CE-968yZTsTLTx1Uj1nM,7004
+compressed_tensors-0.9.5a20250424.dist-info/WHEEL,sha256=SmOxYU7pzNKBqASvQJ7DjX3XGUF92lrGhMb3R6_iiqI,91
+compressed_tensors-0.9.5a20250424.dist-info/top_level.txt,sha256=w2i-GyPs2s1UwVxvutSvN_lM22SXC2hQFBmoMcPnV7Y,19
+compressed_tensors-0.9.5a20250424.dist-info/RECORD,,

{compressed_tensors-0.9.4a20250421.dist-info → compressed_tensors-0.9.5a20250424.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (79.0.0)
+Generator: setuptools (79.0.1)
 Root-Is-Purelib: true
 Tag: py3-none-any