From f3b0cbf998dd767b426ef8c8197f47c73f4727a3 Mon Sep 17 00:00:00 2001
From: William Falcon <falc@fb.com>
Date: Wed, 24 Jul 2019 16:19:19 -0400
Subject: [PATCH] removed dead code in grads

---
 pytorch_lightning/root_module/memory.py      | 13 +++++++++----
 pytorch_lightning/root_module/root_module.py |  1 +
 2 files changed, 10 insertions(+), 4 deletions(-)

diff --git a/pytorch_lightning/root_module/memory.py b/pytorch_lightning/root_module/memory.py
index 0bafe56220..4a166d63a5 100644
--- a/pytorch_lightning/root_module/memory.py
+++ b/pytorch_lightning/root_module/memory.py
@@ -115,13 +115,16 @@ class ModelSummary(object):
         '''
 
         df = pd.DataFrame( np.zeros( (len(self.layer_names), 3) ) )
-        df.columns = ['Name', 'Type', 'Params', 'In_sizes', 'Out_sizes']
+        df.columns = ['Name', 'Type', 'Params']
 
         df['Name'] = self.layer_names
         df['Type'] = self.layer_types
         df['Params'] = self.param_nums
-        df['In_sizes'] = self.in_sizes
-        df['Out_sizes'] = self.out_sizes
+
+        if self.example_input_array:
+            df.columns.extend(['In_sizes', 'Out_sizes'])
+            df['In_sizes'] = self.in_sizes
+            df['Out_sizes'] = self.out_sizes
 
         self.summary = df
         return
@@ -130,7 +133,9 @@ class ModelSummary(object):
         self.get_layer_names()
         self.get_parameter_sizes()
         self.get_parameter_nums()
-        self.get_variable_sizes()
+
+        if self.example_input_array:
+            self.get_variable_sizes()
         self.make_summary()
 
 
diff --git a/pytorch_lightning/root_module/root_module.py b/pytorch_lightning/root_module/root_module.py
index 2345997b00..7f99ef98be 100644
--- a/pytorch_lightning/root_module/root_module.py
+++ b/pytorch_lightning/root_module/root_module.py
@@ -21,6 +21,7 @@ class LightningModule(GradInformation, ModelIO, ModelHooks):
         self.loaded_optimizer_states_dict = {}
         self.trainer = None
         self.experiment = None
+        self.example_input_array = None
 
         # track if gpu was requested for checkpointing
         self.on_gpu = False