Update app.py

app.py

@@ -58,7 +58,7 @@ def next_image():
             )
 
             embs.append(pooled_embeds)
-            return image[0]
+            return image
         else:
             print('######### Roaming #########')
 
@@ -98,7 +98,7 @@ def next_image():
             embs.append(im_emb)
 
             torch.save(lin_class.coef_, f'./{start_time}.pt')
-            return image[0]
+            return image