Rooba
/
lightning
mirror of https://github.com/Lightning-AI/lightning.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

fixed correct module on hpc save

This commit is contained in:
William Falcon 2019-07-24 18:18:58 -04:00
parent 7fa759ffed
commit 8f0d9af168
1 changed files with 3 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
tests/test_models.py
View File

@ -305,6 +305,9 @@ def run_gpu_model_test(trainer_options, model, hparams, on_gpu=True):
# test model preds # test model preds
run_prediction(model.test_dataloader, pretrained_model) run_prediction(model.test_dataloader, pretrained_model)
if trainer.use_ddp:
trainer.model = pretrained_model
# test HPC loading / saving # test HPC loading / saving
trainer.hpc_save(save_dir, exp) trainer.hpc_save(save_dir, exp)
trainer.hpc_load(save_dir, on_gpu=on_gpu) trainer.hpc_load(save_dir, on_gpu=on_gpu)