recompile_limit reached

[AmitMY]

W0110 12:27:06.332000 158 torch/_dynamo/convert_frame.py:1551] [2/8] torch._dynamo hit config.recompile_limit (8)
W0110 12:27:06.332000 158 torch/_dynamo/convert_frame.py:1551] [2/8] function: 'dion2_post_orthogonalize' (/usr/local/lib/python3.12/dist-packages/dion/dion2.py:621)
W0110 12:27:06.332000 158 torch/_dynamo/convert_frame.py:1551] [2/8] last reason: 2/7: tensor 'X[0]' size mismatch at index 1. expected 640, actual 8. Guard failed on a parameter, consider using torch._dynamo.config.force_parameter_static_shapes = False to allow dynamism on parameters.

torch._dynamo.exc.FailOnRecompileLimitHit: recompile_limit reached with fullgraph=True. Excessive recompilations can degrade performance due to the compilation overhead of each recompilation. To monitor recompilations, enable TORCH_LOGS=recompiles. If recompilations are expected, consider increasing torch._dynamo.config.cache_size_limit to an appropriate value.

[AmitMY]

Solution:

import torch._dynamo
# Configure torch.compile to handle dynamic shapes in Dion optimizer
torch._dynamo.config.cache_size_limit = 64
torch._dynamo.config.force_parameter_static_shapes = False