enhance import&add pack ut

Signed-off-by: Kaihui-intel <[email protected]>
intel · Jul 16, 2024 · 0ea77fd · 0ea77fd
1 parent 8dbf793
commit 0ea77fd
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 1 deletion.
diff --git a/neural_compressor/torch/__init__.py b/neural_compressor/torch/__init__.py
@@ -11,3 +11,4 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+from neural_compressor.torch.algorithms.layer_wise import load_empty_model
diff --git a/test/3x/torch/quantization/weight_only/test_rtn.py b/test/3x/torch/quantization/weight_only/test_rtn.py
@@ -167,7 +167,7 @@ def test_quant_lm_head(self):
         ), "The tied lm_head weight is not deep copied, please check!"
 
     def test_layer_wise(self):
-        from neural_compressor.torch.algorithms.layer_wise import load_empty_model
+        from neural_compressor.torch import load_empty_model
 
         model = load_empty_model("hf-internal-testing/tiny-random-GPTJForCausalLM")
         quant_config = RTNConfig(