Fix AOT compilation dynamic_shapes to match expected arg names for torch.export.export

app.py

@@ -495,15 +495,14 @@ def chat_response(user_msg, chat_history, system_prompt,
                 
                 # Define dynamic shapes for variable sequence lengths
                 seq_dim = torch.export.Dim('seq', min=1, max=4096)
-                dynamic_shapes = tree_map(lambda v: None, call.kwargs)
-                
-                # Set dynamic dimensions for common inputs
-                if 'input_ids' in call.kwargs:
-                    dynamic_shapes['input_ids'] = {1: seq_dim}
-                if 'attention_mask' in call.kwargs:
-                    dynamic_shapes['attention_mask'] = {1: seq_dim}
-                if 'position_ids' in call.kwargs:
-                    dynamic_shapes['position_ids'] = {1: seq_dim}
+                dynamic_shapes = {
+                    'input_ids': {1: seq_dim} if 'input_ids' in call.kwargs else None,
+                    'attention_mask': {1: seq_dim} if 'attention_mask' in call.kwargs else None,
+                    'inputs_embeds': None,
+                    'use_cache': None,
+                    'cache_position': {1: seq_dim} if 'cache_position' in call.kwargs or 'position_ids' in call.kwargs else None,
+                    'kwargs': None
+                }
                 
                 exported = torch.export.export(
                     pipe.model,