Fix bug & apply review recommendations

pytorch · Jan 15, 2025 · bbc8dcd · bbc8dcd
1 parent da2e9e0
commit bbc8dcd
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 3 deletions.
diff --git a/test/integration/test_integration.py b/test/integration/test_integration.py
@@ -10,6 +10,7 @@
 import logging
 import os
 import unittest
+from functools import partial
 
 import torch
 import torch.nn as nn
@@ -1004,8 +1005,6 @@ def _test_lin_weight_subclass_api_impl(
     def test_int8_dynamic_quant_subclass_api(
         self, device, dtype, act_mapping, weight_zero_point_domain
     ):
-        from functools import partial
-
         if (
             not TORCH_VERSION_AT_LEAST_2_5
             and dtype in (torch.float16, torch.bfloat16)

diff --git a/torchao/quantization/quant_primitives.py b/torchao/quantization/quant_primitives.py
@@ -372,7 +372,7 @@ def _quantize_affine_no_dtype_cast(
     zero_point: Optional[torch.Tensor],
     quant_min: Union[int, float],
     quant_max: Union[int, float],
-    quant_dtype: Optional[torch.dtype],
+    quant_dtype: torch.dtype,
     zero_point_domain: Optional[str] = ZeroPointDomain.INT.name,
 ) -> torch.Tensor:
     """