Hello author, I tried using the original DiT model for training but facing out of memory issue. I saw your repository which implements DiT using memory constraints. In the README file, I saw you used a mixed_precision argument but I couldn't find it anywhere in the code. I just want to copy the model architecture file and adjust it according to my implementation of the work. Can you please tell which model arch uses less memory constraints as it is a bit confusing to me to understand so just clarifying.