From 2e0fde7da7520beff8c7384e7603e6b8278e224d Mon Sep 17 00:00:00 2001
From: William Falcon <falc@fb.com>
Date: Wed, 24 Jul 2019 18:11:29 -0400
Subject: [PATCH] fixed correct module on hpc save

---
 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 04ca68c65b..ccd63398e6 100644
--- a/tests/test_models.py
+++ b/tests/test_models.py
@@ -277,7 +277,7 @@ def run_gpu_model_test(trainer_options, model, hparams, on_gpu=True):
 
     # test HPC loading / saving
     trainer.hpc_save(save_dir, exp)
-    trainer.hpc_load(save_dir, on_gpu=True)
+    trainer.hpc_load(save_dir, on_gpu=on_gpu)
 
     clear_save_dir()