ptq test error correction (#2860)

pytorch · Jun 7, 2024 · 4f82ea7 · 4f82ea7
1 parent 2645e21
commit 4f82ea7
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 4 deletions.
diff --git a/tests/py/ts/ptq/test_ptq_dataloader_calibrator.py b/tests/py/ts/ptq/test_ptq_dataloader_calibrator.py
@@ -6,6 +6,7 @@
 import torch_tensorrt as torchtrt
 import torchvision
 import torchvision.transforms as transforms
+import torch_tensorrt.ts.ptq as PTQ
 from torch.nn import functional as F
 from torch_tensorrt.ts.logging import *
 
@@ -76,11 +77,11 @@ def test_compile_script(self):
         self.testing_dataloader = torch.utils.data.DataLoader(
             self.testing_dataset, batch_size=1, shuffle=False, num_workers=1
         )
-        self.calibrator = torchtrt.ptq.DataLoaderCalibrator(
+        self.calibrator = PTQ.DataLoaderCalibrator(
             self.testing_dataloader,
             cache_file="./calibration.cache",
             use_cache=False,
-            algo_type=torchtrt.ptq.CalibrationAlgo.ENTROPY_CALIBRATION_2,
+            algo_type=PTQ.CalibrationAlgo.ENTROPY_CALIBRATION_2,
             device=torch.device("cuda:0"),
         )
 

diff --git a/tests/py/ts/ptq/test_ptq_to_backend.py b/tests/py/ts/ptq/test_ptq_to_backend.py
@@ -6,6 +6,7 @@
 import torch_tensorrt as torchtrt
 import torchvision
 import torchvision.transforms as transforms
+import torch_tensorrt.ts.ptq as PTQ
 from torch.nn import functional as F
 from torch_tensorrt.ts.logging import *
 
@@ -76,11 +77,11 @@ def test_compile_script(self):
         self.testing_dataloader = torch.utils.data.DataLoader(
             self.testing_dataset, batch_size=1, shuffle=False, num_workers=1
         )
-        self.calibrator = torchtrt.ptq.DataLoaderCalibrator(
+        self.calibrator = PTQ.DataLoaderCalibrator(
             self.testing_dataloader,
             cache_file="./calibration.cache",
             use_cache=False,
-            algo_type=torchtrt.ptq.CalibrationAlgo.ENTROPY_CALIBRATION_2,
+            algo_type=PTQ.CalibrationAlgo.ENTROPY_CALIBRATION_2,
             device=torch.device("cuda:0"),
         )