Reproducibility results of Sudoku-Extreme-1k

Hey,

Thank you for providing a clean repo for understanding HRM better. I am trying to reproduce the results on Sudoku-Extreme-1k. I have diligently followed all the instructions and other discussions (See https://github.com/sapientinc/HRM/issues/12). However, I am unable to reproduce the results on A100 GPUs.

I think the the underlying issue could be the use of Flash Attention-3 in H100s while I use Flash Attention-2 in a single A100 GPUs. Would it be possible to provide another set of baseline results for HRM on A100 / any other GPUs with falsh attention 2? Please let me know.

This would really help budding researchers with limited compute resources to compare and contrast with HRM.

Thanks in advance,