Update app.py

app.py

@@ -256,7 +256,10 @@ def prod_function(model_pretrained, prod_ds):
 
     #prod_logits = outputs.squeeze(1)
     #prod_pred = prod_logits.argmax(-1)
-    prod_pred = outputs.argmax(-1)
+    
+    prod_logits = outputs.logits
+    prod_softmax = torch.nn.functional.softmax(prod_logits, dim=-1)
+    prod_pred = prod_softmax.argmax(-1)
     return prod_pred
 
 # Function to get landmarked video