Update optimization.py

optimization.py

@@ -8,7 +8,7 @@ from torchao.quantization import quantize_
 from torchao.quantization import Float8DynamicActivationFloat8WeightConfig
 import spaces
 import torch
-from torch.utils._pytree import tree_map
+from torch.utils._pytree import tree_map, tree_map_only
 
 
 P = ParamSpec('P')
@@ -39,10 +39,11 @@ def optimize_pipeline_(pipeline: Callable[P, Any], *args: P.args, **kwargs: P.kw
         with spaces.aoti_capture(pipeline.transformer) as call:
             pipeline(*args, **kwargs)
 
-        dynamic_shapes = tree_map(lambda t: None, call.kwargs)
+        #dynamic_shapes = tree_map(lambda t: None, call.kwargs)
+        dynamic_shapes = tree_map_only((torch.Tensor, bool), lambda x: None, call.kwargs)
         dynamic_shapes |= TRANSFORMER_DYNAMIC_SHAPES
 
-        quantize_(pipeline.transformer, Float8DynamicActivationFloat8WeightConfig())
+        #quantize_(pipeline.transformer, Float8DynamicActivationFloat8WeightConfig())
         
         exported = torch.export.export(
             mod=pipeline.transformer,