kaupane's picture
Update models/DiT.py
23f7c2b verified