Hey, really good stuff! Just one suggestion, possibly maybe change x in your solutions to be more meaningful, up to you tho! Also there's a slight bit of duplication of code where the forward function is concerned. Id advise making like a base functions class. feel free to ignore me tho its still pretty good!
#PyTorch Implementation of T5
fading ridge
keen tendon
Pushed a final update. I think it is pretty close to the original in the paper now.