Update main.py

main.py

@@ -8,7 +8,7 @@ import timm
 import torch
 import torch.nn.functional as F
 
-device = torch.device("cpu")
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 class CFG:
     debug = False
@@ -296,7 +296,7 @@ def get_image_embeddings(image):
     return image_embeddings
 
 
-def predict_caption(image, model, text_embeddings, captions, n=2):
+def predict_caption(image, model, text_embeddings, captions, n=1):
     # get the image embeddings
     image_embeddings = get_image_embeddings(image)
     if image_embeddings is None:
@@ -332,4 +332,4 @@ def get_text_embeddings(valid_df):
             text_embeddings = model.text_projection(text_features)
             valid_text_embeddings.append(text_embeddings)
 
-    return model, torch.cat(valid_text_embeddings)
+    return model, torch.cat(valid_text_embeddings)