augmented utils.py to account for 100% training set