Updated training functions

Training functions now train based on two classes. Test functions will
be next to do the same.