updated to fix the None problem