Fixing tests

Deci-AI · Apr 19, 2024 · 8be4575 · 8be4575
1 parent 0a3c076
commit 8be4575
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 2 deletions.
diff --git a/src/super_gradients/training/sg_trainer/sg_trainer.py b/src/super_gradients/training/sg_trainer/sg_trainer.py
@@ -2778,7 +2778,10 @@ def _export_quantized_model(
 
             device = "cpu"
             example_input_shape = next(iter(quantization_result.calibration_dataloader))[0].size()
-            input_shape_with_explicit_batch = tuple(export_params.batch_size, *example_input_shape[1:])
+            input_shape_with_explicit_batch = tuple([export_params.batch_size, *example_input_shape[1:]])
+            if export_params.input_image_shape is not None:
+                input_shape_with_explicit_batch = input_shape_with_explicit_batch[: -len(export_params.input_image_shape)] + export_params.input_image_shape
+
             onnx_input = torch.randn(input_shape_with_explicit_batch).to(device=device)
             onnx_export_kwargs = export_params.onnx_export_kwargs or {}
             model_to_export = quantization_result.quantized_model

diff --git a/tests/recipe_training_tests/coded_qat_launch_test.py b/tests/recipe_training_tests/coded_qat_launch_test.py
@@ -1,13 +1,18 @@
 import unittest
 
+from super_gradients.import_utils import import_pytorch_quantization_or_install
 from torchvision.transforms import Normalize, ToTensor, RandomHorizontalFlip, RandomCrop
 
 from super_gradients import Trainer
-from super_gradients.training import modify_params_for_qat
+
 from super_gradients.training.dataloaders.dataloaders import cifar10_train, cifar10_val
 from super_gradients.training.metrics import Accuracy, Top5
 from super_gradients.training.models import ResNet18
 
+from super_gradients.training.utils.quantization.tensorrt.functional import modify_params_for_qat
+
+import_pytorch_quantization_or_install()
+
 
 class CodedQATLuanchTest(unittest.TestCase):
     def test_qat_launch(self):