Update optimization.py

optimization.py

@@ -7,13 +7,13 @@ from typing import ParamSpec
 
 import spaces
 import torch
-from torch.utils._pytree import tree_map_only
+from torch.utils._pytree import tree_map
 
 
 P = ParamSpec('P')
 
 
-TRANSFORMER_HIDDEN_DIM = torch.export.Dim('hidden', min=3584, max=9727)
+TRANSFORMER_HIDDEN_DIM = torch.export.Dim.AUTO('hidden', min=3584, max=9727)
 
 TRANSFORMER_DYNAMIC_SHAPES = {
     'hidden_states': {1: TRANSFORMER_HIDDEN_DIM},
@@ -38,7 +38,7 @@ 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_only((torch.Tensor, bool, int), lambda t: None, call.kwargs)
+        dynamic_shapes = tree_map(lambda t: None, call.kwargs)
         dynamic_shapes |= TRANSFORMER_DYNAMIC_SHAPES
         
         exported = torch.export.export(