Update optimization.py

optimization.py

@@ -39,8 +39,8 @@ 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_only((torch.Tensor, bool), lambda x: 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())