Update main.py

main.py

@@ -20,9 +20,9 @@ def under_max(image):
   return image
 
 class Model(object):
-  def __init__(self, gpu=0):
+  def __init__(self, gpu=None):
     config = Config()
-    config.device = 'cuda:{}'.format(gpu)
+    config.device = 'cpu' if gpu is None else 'cuda:{}'.format(gpu)
     model, _ = caption_model.build_model(config)   
     checkpoint = torch.load('./checkpoint.pth', map_location='cpu')
     model.load_state_dict(checkpoint['model'])