Minor fixes for bugs that appeared after optimization