Create an Enum for tracking Train, validation, test stage 
🚀 Feature

Motivation
The goal is to create single point of truth about stages for the trainer, so we don't get typo errors.

Pitch

Alternatives

Additional context