Fix tests

src/brevitas_examples/imagenet_classification/qat/__init__.py

@@ -1,5 +1,2 @@
 # Copyright (C) 2023, Advanced Micro Devices, Inc. All rights reserved.
 # SPDX-License-Identifier: BSD-3-Clause
-
-
-from .models import *

src/brevitas_examples/imagenet_classification/quant_models/__init__.py

@@ -0,0 +1,4 @@
+# Copyright (C) 2023, Advanced Micro Devices, Inc. All rights reserved.
+# SPDX-License-Identifier: BSD-3-Clause
+
+from .models import *

tests/brevitas/export/test_generic_export.py

@@ -15,7 +15,7 @@
 from brevitas.quant.scaled_int import Int4WeightPerTensorFloatDecoupled
 from brevitas.quant.scaled_int import Int8ActPerTensorFloat
 from brevitas.quant.scaled_int import Int16Bias
-from brevitas_examples import imagenet_classification
+from brevitas_examples.imagenet_classification import quant_models
 from tests.marker import jit_disabled_for_export
 
 OUT_CH = 50
@@ -173,7 +173,7 @@ def test_generic_quant_avgpool_export_quant_input():
 
 @jit_disabled_for_export()
 def test_debug_brevitas_onnx_export():
- model, cfg = imagenet_classification.model_with_cfg('quant_mobilenet_v1_4b', pretrained=False)
+ model, cfg = quant_models.model_with_cfg('quant_mobilenet_v1_4b', pretrained=False)
  model.eval()
  debug_hook = enable_debug(model, proxy_level=True)
  input_tensor = torch.randn(1, 3, 224, 224)

tests/brevitas_examples/test_examples_import.py

@@ -17,10 +17,13 @@ def test_import_bnn_pynq():
 
 
 def test_import_image_classification():
- from brevitas_examples.imagenet_classification import quant_mobilenet_v1_4b
- from brevitas_examples.imagenet_classification import quant_proxylessnas_mobile14_4b
- from brevitas_examples.imagenet_classification import quant_proxylessnas_mobile14_4b5b
- from brevitas_examples.imagenet_classification import quant_proxylessnas_mobile14_hadamard_4b
+ from brevitas_examples.imagenet_classification.quant_models import quant_mobilenet_v1_4b
+ from brevitas_examples.imagenet_classification.quant_models import \
+ quant_proxylessnas_mobile14_4b
+ from brevitas_examples.imagenet_classification.quant_models import \
+ quant_proxylessnas_mobile14_4b5b
+ from brevitas_examples.imagenet_classification.quant_models import \
+ quant_proxylessnas_mobile14_hadamard_4b
 
 
 def test_import_tts():

tests/brevitas_finn/brevitas_examples/test_mobilenet_finn_export.py

@@ -17,7 +17,7 @@
 
 from brevitas import torch_version
 from brevitas.export import export_finn_onnx
-from brevitas_examples.imagenet_classification import quant_mobilenet_v1_4b
+from brevitas_examples.imagenet_classification.quant_models import quant_mobilenet_v1_4b
 
 ort_mac_fail = pytest.mark.skipif(
  torch_version >= parse('1.5.0') and system() == 'Darwin',