From b824f184ff282cc316a7165a1a8b23f43c13869a Mon Sep 17 00:00:00 2001 From: William Falcon Date: Wed, 24 Jul 2019 16:31:55 -0400 Subject: [PATCH] added sample input for summary --- pytorch_lightning/root_module/memory.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pytorch_lightning/root_module/memory.py b/pytorch_lightning/root_module/memory.py index 29d6e66d..6f29b9ad 100644 --- a/pytorch_lightning/root_module/memory.py +++ b/pytorch_lightning/root_module/memory.py @@ -123,11 +123,11 @@ class ModelSummary(object): Layer Name, Layer Type, Input Size, Output Size, Number of Parameters ''' - df = pd.DataFrame( np.zeros( (len(self.layer_names), 3) ) ) cols = ['Name', 'Type', 'Params'] if self.model.example_input_array is not None: cols.extend(['In_sizes', 'Out_sizes']) + df = pd.DataFrame(np.zeros( (len(self.layer_names), len(cols)))) df.columns = cols df['Name'] = self.layer_names