Small patch to fix last problem