From 7f420c0cc2bd727a6b139dde2bb25b48ddfe13f3 Mon Sep 17 00:00:00 2001
From: William Falcon <falc@fb.com>
Date: Wed, 24 Jul 2019 17:41:08 -0400
Subject: [PATCH] test memory printing

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

diff --git a/tests/test_models.py b/tests/test_models.py
index 4195f9a6..42c59738 100644
--- a/tests/test_models.py
+++ b/tests/test_models.py
@@ -6,6 +6,7 @@ from argparse import Namespace
 from test_tube import Experiment
 from pytorch_lightning.callbacks import ModelCheckpoint, EarlyStopping
 from pytorch_lightning.utils.debugging import MisconfigurationException
+from pytorch_lightning.root_module import memory
 import numpy as np
 import warnings
 import torch
@@ -128,6 +129,11 @@ def test_multi_gpu_model_dp():
 
     run_gpu_model_test(trainer_options, model, hparams)
 
+    # test memory helper functions
+    memory.count_mem_items()
+    memory.print_mem_stack()
+    memory.get_gpu_memory_map()
+
 
 def test_amp_gpu_dp():
     """
@@ -234,14 +240,12 @@ def test_ddp_sampler_error():
         use_amp=True
     )
 
-    # test memory gathering
-    trainer.count_mem_items()
-
     with pytest.raises(MisconfigurationException):
         trainer.get_dataloaders(model)
 
     clear_save_dir()
 
+
 # ------------------------------------------------------------------------
 # UTILS
 # ------------------------------------------------------------------------