Update optimization.py

optimization.py

@@ -125,6 +125,8 @@ def optimize_pipeline_(pipeline: Callable[P, Any], *args: P.args, **kwargs: P.kw
 
     transformer_config = pipeline.transformer.config
     transformer_dtype = pipeline.transformer.dtype
+    cache_context = pipeline.transformer.cache_context
     pipeline.transformer = combined_transformer
     pipeline.transformer.config = transformer_config  # pyright: ignore[reportAttributeAccessIssue]
     pipeline.transformer.dtype = transformer_dtype  # pyright: ignore[reportAttributeAccessIssue]
+    pipeline.transformer.cache_context = cache_context