Update

[ghstack-poisoned]

Author: vkuzo

test/dtypes/test_affine_quantized.py

@@ -189,7 +189,7 @@ def test_flatten_unflatten(self, device, dtype):
         for apply_quant in apply_quant_list:
             linear = torch.nn.Linear(128, 256, dtype=dtype, device=device)
             if isinstance(apply_quant, AOBaseWorkflowConfig):
-                quantize_(linear, apply_quant) 
+                quantize_(linear, apply_quant)
             else:
                 ql = apply_quant(linear)
             lp_tensor = ql.weight

torchao/quantization/quant_api.py

@@ -481,9 +481,7 @@ def insert_subclass(lin):
 
 def quantize_(
     model: torch.nn.Module,
-    config: Union[
-        AOBaseWorkflowConfig, Callable[[torch.nn.Module], torch.nn.Module]
-    ],
+    config: Union[AOBaseWorkflowConfig, Callable[[torch.nn.Module], torch.nn.Module]],
     filter_fn: Optional[Callable[[torch.nn.Module, str], bool]] = None,
     set_inductor_config: bool = True,
     device: Optional[torch.types.Device] = None,
@@ -535,7 +533,9 @@ def quantize_(
 
     else:
         # old behavior, keep to avoid breaking BC
-        warnings.warn("""Passing a generic Callable to `quantize_` is no longer recommended and will be deprecated at a later release. Please see https://github.com/pytorch/ao/pull/1595 for instructions on how to pass in workflow configuration instead.""")
+        warnings.warn(
+            """Passing a generic Callable to `quantize_` is no longer recommended and will be deprecated at a later release. Please see https://github.com/pytorch/ao/pull/1595 for instructions on how to pass in workflow configuration instead."""
+        )
 
         # make the variable name make sense
         apply_tensor_subclass = config