Fixed up some code