Added training logic
Cleaned up