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

---
 tests/test_models.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/tests/test_models.py b/tests/test_models.py
index a559c339..3a364aea 100644
--- a/tests/test_models.py
+++ b/tests/test_models.py
@@ -305,6 +305,9 @@ def run_gpu_model_test(trainer_options, model, hparams, on_gpu=True):
     # test model preds
     run_prediction(model.test_dataloader, pretrained_model)
 
+    if trainer.use_ddp:
+        trainer.model = pretrained_model
+
     # test HPC loading / saving
     trainer.hpc_save(save_dir, exp)
     trainer.hpc_load(save_dir, on_gpu=on_gpu)