Update app.py

app.py

@@ -54,6 +54,7 @@ class CFG():
     num_workers=1
 
 if st.button('predict'):
+    st.progress(0)
     
     device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
     
@@ -127,12 +128,13 @@ if st.button('predict'):
         preds = []
         model.eval()
         model.to(device)
-        tk0 = tqdm(test_loader, total=len(test_loader))
-        for inputs in tk0:
+        tk0 = enumerate(test_loader, total=len(test_loader))
+        for i, inputs in tk0:
             for k, v in inputs.items():
                 inputs[k] = v.to(device)
             with torch.no_grad():
                 y_preds = model(inputs)
+            st.progress((i+1)*CFG.batch_size)
             preds.append(y_preds.to('cpu').numpy())
         predictions = np.concatenate(preds)
         return predictions
@@ -165,10 +167,11 @@ if st.button('predict'):
         )
     
     else:
+        CFG.batch_size=1
         test_ds = pd.DataFrame.from_dict({'input': CFG.data}, orient='index').T
         test_dataset = TestDataset(CFG, test_ds)
         test_loader = DataLoader(test_dataset,
-                                 batch_size=1,
+                                 batch_size=CFG.batch_size,
                                  shuffle=False,
                                  num_workers=CFG.num_workers, pin_memory=True, drop_last=False)