fixed minor bug in if __main__ section
(#297)