.

Author: Borda

tests/helpers/pipelines.py

@@ -105,21 +105,20 @@ def run_model_test(
         trainer.checkpoint_connector.hpc_load(checkpoint_path, on_gpu=on_gpu)
 
 
+@torch.no_grad()
 def run_prediction_eval_model_template(trained_model, dataloader, dp=False, min_acc=0.50):
     # run prediction on 1 batch
     batch = next(iter(dataloader))
     x, y = batch
     x = x.view(x.size(0), -1)
 
     if dp:
-        with torch.no_grad():
-            output = trained_model(batch, 0)
-            acc = output['val_acc']
+        output = trained_model(batch, 0)
+        acc = output['val_acc']
         acc = torch.mean(acc).item()
 
     else:
-        with torch.no_grad():
-            y_hat = trained_model(x)
-        acc = accuracy(y_hat.cpu(), y.cpu()).item()
+        y_hat = trained_model(x)
+        acc = accuracy(y_hat.cpu(), y.cpu(), top_k=2).item()
 
     assert acc >= min_acc, f"This model is expected to get > {min_acc} in test set (it got {acc})"