ensembles

main.py

@@ -118,11 +118,11 @@ def run_lightning_loop(config_obj):
 
                     from tqdm import tqdm
                     for batch in tqdm(test_dataloader, total=len(test_dataloader)):
-                        batch_x, file_name = batch
+                        batch_x, file_names = batch
                         batch_x = batch_x.to(device='cuda' if model.on_gpu else 'cpu')
                         y = model(batch_x).main_out
                         predictions = (y >= 0.5).int()
-                        for prediction in predictions:
+                        for prediction, file_name in zip(predictions, file_names):
                             prediction_text = 'clear' if prediction == V.CLEAR else 'mask'
                             outfile.write(f'{file_name},{prediction_text}\n')
     return model