Update optimization.py

optimization.py

@@ -1,6 +1,21 @@
 import torch
 import torchao
 from torchao.quantization import DEFAULT_INT4_AUTOQUANT_CLASS_LIST
+from typing import Any
+from typing import Callable
+from typing import ParamSpec
+
+import spaces
+import torch
+from torch.utils._pytree import tree_map_only
+from torchao.quantization import quantize_
+from torchao.quantization import Float8DynamicActivationFloat8WeightConfig
+from torchao.quantization import Int8WeightOnlyConfig
+
+from optimization_utils import capture_component_call
+from optimization_utils import aoti_compile
+from optimization_utils import ZeroGPUCompiledModel
+from optimization_utils import drain_module_parameters
 P = ParamSpec('P')