extend dimension to include batch_size

app.py

@@ -240,7 +240,7 @@ def prod_function(model_pretrained, prod_ds):
     accelerated_model.eval()
 
     with torch.no_grad():
-        outputs = accelerated_model(acclerated_prod_ds)
+        outputs = accelerated_model(acclerated_prod_ds.unsqueeze(0))
 
     prod_logits = outputs.squeeze(1)
     prod_pred = prod_logits.argmax(-1)