Minor code revision