Update optimization.py

optimization.py

@@ -4,7 +4,7 @@
 from typing import Any
 from typing import Callable
 from typing import ParamSpec
-
+from torchao.quantization import Float8DynamicActivationFloat8WeightConfig
 import spaces
 import torch
 from torch.utils._pytree import tree_map
@@ -40,6 +40,8 @@ def optimize_pipeline_(pipeline: Callable[P, Any], *args: P.args, **kwargs: P.kw
 
         dynamic_shapes = tree_map(lambda t: None, call.kwargs)
         dynamic_shapes |= TRANSFORMER_DYNAMIC_SHAPES
+
+        quantize_(pipeline.transformer, Float8DynamicActivationFloat8WeightConfig())
         
         exported = torch.export.export(
             mod=pipeline.transformer,