From a6ae97ac0922382a1e2ed7c593bb064f141001a3 Mon Sep 17 00:00:00 2001
From: William Falcon <falc@fb.com>
Date: Fri, 26 Jul 2019 22:13:06 -0400
Subject: [PATCH] fixed hpc save, load. cleaned apu

---
 tests/test_models.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/test_models.py b/tests/test_models.py
index 1a3156250d..bdd31a5c75 100644
--- a/tests/test_models.py
+++ b/tests/test_models.py
@@ -115,7 +115,7 @@ def test_model_saving_loading():
     # make prediction
     # assert that both predictions are the same
     new_pred = model_2(x)
-    assert torch.eq(pred_before_saving, new_pred)
+    assert torch.all(torch.eq(pred_before_saving, new_pred)).item() == 1
 
     clear_save_dir()