minor change in mnist function