Minor fix to python UI