'NoneType' fix