Some minor fixes to main.go