Update main.py

main.py

@@ -17,8 +17,9 @@ def initialize():
         T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
     ]) 
 
-    checkpoint_path = "checkpoint_epoch_40.pth"
-    model = ImgCap(cnn_feature_size=1024, lstm_hidden_size=1024, embedding_dim=1024, num_layers=2, vocab_size=len(vocab))
+    checkpoint_path = "checkpoint_epoch_30.pth"
+    model = ImgCap(feature_size=2048, lstm_hidden_size=1024, embedding_dim=1024, num_layers=2, vocab_size=len(vocab))
+    
     model, _, _, _, _, _, _ = load_checkpoint(checkpoint_path=checkpoint_path, model=model)
 
     return model, vocab, transforms