I swear I did this before. Also, this one is garbage because Colab cut off my GPU for the day.

initial commit