Added new parameter, so we can train on specific classes, instead of having to actually test everything out.