From 5f814e48c4162a50bd90bda65191a372594a9526 Mon Sep 17 00:00:00 2001 From: William Falcon Date: Wed, 24 Jul 2019 16:30:27 -0400 Subject: [PATCH] added sample input for summary --- pytorch_lightning/root_module/memory.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/pytorch_lightning/root_module/memory.py b/pytorch_lightning/root_module/memory.py index a2feb689..29d6e66d 100644 --- a/pytorch_lightning/root_module/memory.py +++ b/pytorch_lightning/root_module/memory.py @@ -124,14 +124,18 @@ class ModelSummary(object): ''' df = pd.DataFrame( np.zeros( (len(self.layer_names), 3) ) ) - df.columns = ['Name', 'Type', 'Params'] + cols = ['Name', 'Type', 'Params'] + if self.model.example_input_array is not None: + cols.extend(['In_sizes', 'Out_sizes']) + + df.columns = cols df['Name'] = self.layer_names df['Type'] = self.layer_types df['Params'] = self.param_nums if self.model.example_input_array is not None: - df.columns.extend(['In_sizes', 'Out_sizes']) + df['In_sizes'] = self.in_sizes df['Out_sizes'] = self.out_sizes