Upgrading the PyTorch Version

[JavadZohani]

Hi

Hope you're doing well, and thank you for the great work on this script.

I wanted to suggest upgrading the PyTorch version used in the project to torch==2.7.1. This version has proven to be more stable and compatible with newer dependencies like CUDA and torchvision in my experience. It also seems to be the recommended version for several other similar projects.

If possible, please consider updating to this version in a future release. It would really help improve stability for many users.

Thanks again for your efforts!

Best regards,

[bmaltais]

Have you tried the dev branch with CUDA 12.8?