compressed-tensors 0.10.2a20250617__tar.gz → 0.10.2a20250620__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (144) hide show
  1. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/.github/workflows/report.yml +1 -1
  2. {compressed_tensors-0.10.2a20250617/src/compressed_tensors.egg-info → compressed_tensors-0.10.2a20250620}/PKG-INFO +1 -1
  3. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/src/compressed_tensors/utils/offload.py +10 -10
  4. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/src/compressed_tensors/version.py +1 -1
  5. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620/src/compressed_tensors.egg-info}/PKG-INFO +1 -1
  6. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/tests/test_utils/test_offload.py +19 -0
  7. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/.github/.gitkeep +0 -0
  8. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/.github/actions/test/action.yml +0 -0
  9. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/.github/scripts/step-status +0 -0
  10. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/.github/workflows/build-test.yml +0 -0
  11. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/.github/workflows/build.yml +0 -0
  12. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/.github/workflows/test-check.yaml +0 -0
  13. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/.github/workflows/test.yml +0 -0
  14. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/.github/workflows/trigger-all.yml +0 -0
  15. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/.github/workflows/upload.yml +0 -0
  16. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/.gitignore +0 -0
  17. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/LICENSE +0 -0
  18. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/Makefile +0 -0
  19. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/README.md +0 -0
  20. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/examples/bit_packing/ex_quantize_and_pack.py +0 -0
  21. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/examples/bit_packing/int4_config.json +0 -0
  22. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/examples/bitmask_compression.ipynb +0 -0
  23. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/examples/llama_1.1b/ex_config_quantization.py +0 -0
  24. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/examples/llama_1.1b/ex_llmcompressor_quantization.py +0 -0
  25. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/examples/llama_1.1b/example_quant_config.json +0 -0
  26. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/examples/llama_1.1b/example_quant_recipe.yaml +0 -0
  27. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/examples/quantize_and_pack_int4.ipynb +0 -0
  28. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/pyproject.toml +0 -0
  29. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/setup.cfg +0 -0
  30. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/setup.py +0 -0
  31. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/src/__init__.py +0 -0
  32. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/src/compressed_tensors/README.md +0 -0
  33. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/src/compressed_tensors/__init__.py +0 -0
  34. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/src/compressed_tensors/base.py +0 -0
  35. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/src/compressed_tensors/compressors/__init__.py +0 -0
  36. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/src/compressed_tensors/compressors/base.py +0 -0
  37. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/src/compressed_tensors/compressors/helpers.py +0 -0
  38. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/src/compressed_tensors/compressors/model_compressors/__init__.py +0 -0
  39. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/src/compressed_tensors/compressors/model_compressors/model_compressor.py +0 -0
  40. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/src/compressed_tensors/compressors/quantized_compressors/__init__.py +0 -0
  41. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/src/compressed_tensors/compressors/quantized_compressors/base.py +0 -0
  42. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/src/compressed_tensors/compressors/quantized_compressors/naive_quantized.py +0 -0
  43. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/src/compressed_tensors/compressors/quantized_compressors/nvfp4_quantized.py +0 -0
  44. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/src/compressed_tensors/compressors/quantized_compressors/pack_quantized.py +0 -0
  45. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/src/compressed_tensors/compressors/sparse_compressors/__init__.py +0 -0
  46. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/src/compressed_tensors/compressors/sparse_compressors/base.py +0 -0
  47. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/src/compressed_tensors/compressors/sparse_compressors/dense.py +0 -0
  48. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/src/compressed_tensors/compressors/sparse_compressors/sparse_24_bitmask.py +0 -0
  49. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/src/compressed_tensors/compressors/sparse_compressors/sparse_bitmask.py +0 -0
  50. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/src/compressed_tensors/compressors/sparse_quantized_compressors/__init__.py +0 -0
  51. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/src/compressed_tensors/compressors/sparse_quantized_compressors/marlin_24.py +0 -0
  52. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/src/compressed_tensors/config/__init__.py +0 -0
  53. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/src/compressed_tensors/config/base.py +0 -0
  54. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/src/compressed_tensors/config/dense.py +0 -0
  55. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/src/compressed_tensors/config/sparse_24_bitmask.py +0 -0
  56. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/src/compressed_tensors/config/sparse_bitmask.py +0 -0
  57. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/src/compressed_tensors/linear/__init__.py +0 -0
  58. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/src/compressed_tensors/linear/compressed_linear.py +0 -0
  59. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/src/compressed_tensors/quantization/__init__.py +0 -0
  60. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/src/compressed_tensors/quantization/lifecycle/__init__.py +0 -0
  61. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/src/compressed_tensors/quantization/lifecycle/apply.py +0 -0
  62. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/src/compressed_tensors/quantization/lifecycle/compressed.py +0 -0
  63. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/src/compressed_tensors/quantization/lifecycle/forward.py +0 -0
  64. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/src/compressed_tensors/quantization/lifecycle/helpers.py +0 -0
  65. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/src/compressed_tensors/quantization/lifecycle/initialize.py +0 -0
  66. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/src/compressed_tensors/quantization/quant_args.py +0 -0
  67. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/src/compressed_tensors/quantization/quant_config.py +0 -0
  68. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/src/compressed_tensors/quantization/quant_scheme.py +0 -0
  69. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/src/compressed_tensors/quantization/utils/__init__.py +0 -0
  70. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/src/compressed_tensors/quantization/utils/helpers.py +0 -0
  71. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/src/compressed_tensors/registry/__init__.py +0 -0
  72. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/src/compressed_tensors/registry/registry.py +0 -0
  73. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/src/compressed_tensors/transform/__init__.py +0 -0
  74. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/src/compressed_tensors/transform/factory/__init__.py +0 -0
  75. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/src/compressed_tensors/transform/factory/base.py +0 -0
  76. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/src/compressed_tensors/transform/factory/hadamard.py +0 -0
  77. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/src/compressed_tensors/transform/factory/matrix_multiply.py +0 -0
  78. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/src/compressed_tensors/transform/factory/random_hadamard.py +0 -0
  79. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/src/compressed_tensors/transform/transform_args.py +0 -0
  80. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/src/compressed_tensors/transform/transform_config.py +0 -0
  81. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/src/compressed_tensors/transform/transform_scheme.py +0 -0
  82. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/src/compressed_tensors/transform/utils/__init__.py +0 -0
  83. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/src/compressed_tensors/transform/utils/hadamard.py +0 -0
  84. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/src/compressed_tensors/transform/utils/hadamards.safetensors +0 -0
  85. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/src/compressed_tensors/transform/utils/utils.py +0 -0
  86. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/src/compressed_tensors/utils/__init__.py +0 -0
  87. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/src/compressed_tensors/utils/helpers.py +0 -0
  88. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/src/compressed_tensors/utils/permutations_24.py +0 -0
  89. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/src/compressed_tensors/utils/permute.py +0 -0
  90. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/src/compressed_tensors/utils/safetensors_load.py +0 -0
  91. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/src/compressed_tensors/utils/semi_structured_conversions.py +0 -0
  92. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/src/compressed_tensors.egg-info/SOURCES.txt +0 -0
  93. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/src/compressed_tensors.egg-info/dependency_links.txt +0 -0
  94. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/src/compressed_tensors.egg-info/requires.txt +0 -0
  95. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/src/compressed_tensors.egg-info/top_level.txt +0 -0
  96. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/tests/__init__.py +0 -0
  97. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/tests/conftest.py +0 -0
  98. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/tests/test_compressors/__init__.py +0 -0
  99. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/tests/test_compressors/model_compressors/__init__.py +0 -0
  100. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/tests/test_compressors/model_compressors/test_model_compressor.py +0 -0
  101. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/tests/test_compressors/quantized_compressors/__init__.py +0 -0
  102. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/tests/test_compressors/quantized_compressors/test_fp8_quant.py +0 -0
  103. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/tests/test_compressors/quantized_compressors/test_int_quant.py +0 -0
  104. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/tests/test_compressors/quantized_compressors/test_nvfp4_quant.py +0 -0
  105. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/tests/test_compressors/quantized_compressors/test_pack_quant.py +0 -0
  106. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/tests/test_compressors/sparse_compressors/__init__.py +0 -0
  107. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/tests/test_compressors/sparse_compressors/test_bitmask.py +0 -0
  108. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/tests/test_compressors/sparse_compressors/test_sparse_24_bitmask.py +0 -0
  109. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/tests/test_compressors/sparse_quantized_compressors/__init__.py +0 -0
  110. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/tests/test_compressors/sparse_quantized_compressors/test_marlin_24.py +0 -0
  111. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/tests/test_configs/__init__.py +0 -0
  112. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/tests/test_configs/test_base.py +0 -0
  113. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/tests/test_examples/test_bitmask_compression_ipynb.py +0 -0
  114. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/tests/test_linear/__init__.py +0 -0
  115. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/tests/test_linear/test_compressed_linear.py +0 -0
  116. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/tests/test_quantization/__init__.py +0 -0
  117. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/tests/test_quantization/lifecycle/__init__.py +0 -0
  118. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/tests/test_quantization/lifecycle/conftest.py +0 -0
  119. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/tests/test_quantization/lifecycle/test_apply.py +0 -0
  120. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/tests/test_quantization/lifecycle/test_dynamic_lifecycle.py +0 -0
  121. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/tests/test_quantization/lifecycle/test_enabled.py +0 -0
  122. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/tests/test_quantization/lifecycle/test_forward.py +0 -0
  123. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/tests/test_quantization/lifecycle/test_helpers.py +0 -0
  124. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/tests/test_quantization/lifecycle/test_initialize.py +0 -0
  125. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/tests/test_quantization/lifecycle/test_lifecycle.py +0 -0
  126. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/tests/test_quantization/test_configs/__init__.py +0 -0
  127. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/tests/test_quantization/test_configs/test_bit_depths.py +0 -0
  128. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/tests/test_quantization/test_configs/test_strategies.py +0 -0
  129. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/tests/test_quantization/test_quant_args.py +0 -0
  130. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/tests/test_quantization/test_quant_config.py +0 -0
  131. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/tests/test_quantization/test_quant_scheme.py +0 -0
  132. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/tests/test_quantization/test_utils/test_helpers.py +0 -0
  133. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/tests/test_registry.py +0 -0
  134. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/tests/test_transform/factory/test_correctness.py +0 -0
  135. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/tests/test_transform/factory/test_memory.py +0 -0
  136. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/tests/test_transform/test_transform_args.py +0 -0
  137. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/tests/test_transform/test_transform_config.py +0 -0
  138. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/tests/test_transform/test_transform_scheme.py +0 -0
  139. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/tests/test_transform/utils/test_hadamard.py +0 -0
  140. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/tests/test_utils/__init__.py +0 -0
  141. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/tests/test_utils/test_helpers.py +0 -0
  142. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/tests/test_utils/test_safetensors_load.py +0 -0
  143. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/tests/testing_utils.py +0 -0
  144. {compressed_tensors-0.10.2a20250617 → compressed_tensors-0.10.2a20250620}/utils/copyright.py +0 -0
@@ -120,7 +120,7 @@ jobs:
120
120
  shell: bash
121
121
 
122
122
  - name: report to reportportal
123
- uses: neuralmagic/nm-actions/actions/reportportal_submit_execution_results@v1.15.0
123
+ uses: neuralmagic/nm-actions/actions/reportportal_submit_execution_results@v1.22.0
124
124
  with:
125
125
  droute_username: ${{ secrets.DROUTE_USERNAME }}
126
126
  droute_password: ${{ secrets.DROUTE_PASSWORD }}
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: compressed-tensors
3
- Version: 0.10.2a20250617
3
+ Version: 0.10.2a20250620
4
4
  Summary: Library for utilization of compressed safetensors of neural network models
5
5
  Home-page: https://github.com/neuralmagic/compressed-tensors
6
6
  Author: Neuralmagic, Inc.
@@ -171,22 +171,22 @@ def update_parameter_data(
171
171
 
172
172
  def get_execution_device(module: torch.nn.Module) -> torch.device:
173
173
  """
174
- Get the device which inputs should be moved to before module execution
174
+ Get the device which inputs should be moved to before module execution.
175
+ Assume that modules execute in the same order as returned by `model.modules()`
175
176
 
176
177
  :param module: module to check, may be offloaded
177
178
  :return: onload device of module
178
179
  """
179
- if has_offloaded_params(module):
180
- return module._hf_hook.execution_device
180
+ for submodule in module.modules():
181
+ if has_offloaded_params(submodule):
182
+ return submodule._hf_hook.execution_device
181
183
 
182
- first_param = next(module.parameters(), None)
183
- if first_param is None:
184
- warnings.warn(
185
- f"Unable able to infer execution device of {module}, falling back to CPU"
186
- )
187
- return torch.device("cpu")
184
+ param = next(submodule.parameters(recurse=False), None)
185
+ if param is not None:
186
+ return param.device
188
187
 
189
- return first_param.device
188
+ warnings.warn(f"Unable to get execution device of {module}, falling back to CPU")
189
+ return torch.device("cpu")
190
190
 
191
191
 
192
192
  def register_offload_parameter(
@@ -17,5 +17,5 @@ __version__: str
17
17
  __version_tuple__: VERSION_TUPLE
18
18
  version_tuple: VERSION_TUPLE
19
19
 
20
- __version__ = version = '0.10.2.a20250617'
20
+ __version__ = version = '0.10.2.a20250620'
21
21
  __version_tuple__ = version_tuple = (0, 10, 2)
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: compressed-tensors
3
- Version: 0.10.2a20250617
3
+ Version: 0.10.2a20250620
4
4
  Summary: Library for utilization of compressed safetensors of neural network models
5
5
  Home-page: https://github.com/neuralmagic/compressed-tensors
6
6
  Author: Neuralmagic, Inc.
@@ -102,6 +102,25 @@ def test_get_execution_device():
102
102
  assert get_execution_device(module) == torch.device("cuda:0")
103
103
 
104
104
 
105
+ @requires_gpu
106
+ @requires_accelerate()
107
+ def test_get_execution_device_model():
108
+ class Model(torch.nn.Module):
109
+ def __init__(self):
110
+ super().__init__()
111
+ self.a = torch.nn.Linear(1, 2)
112
+ self.b = torch.nn.Linear(2, 2, device="cuda:0")
113
+
114
+ def forward(self, x):
115
+ return self.b(self.a(x).to("cuda:0"))
116
+
117
+ model = Model()
118
+ assert get_execution_device(model) == torch.device("cpu")
119
+
120
+ offloaded_dispatch(model.a, torch.device("cuda:0"))
121
+ assert get_execution_device(model) == torch.device("cuda:0")
122
+
123
+
105
124
  @requires_accelerate()
106
125
  def test_register_offload_parameter():
107
126
  from accelerate import init_empty_weights