Fix ml bug