Using built-in optim
Started grad functoin
Merged optim use
Added train/val/test import