Update main.go

fix bug