added sample input for summary
This commit is contained in:
parent
f3b0cbf998
commit
83ccd21bec
|
@ -28,6 +28,8 @@ class LightningTemplateModel(LightningModule):
|
|||
|
||||
self.batch_size = hparams.batch_size
|
||||
|
||||
self.example_input_array = torch.rand(5, 3 * 28 * 28)
|
||||
|
||||
# build model
|
||||
self.__build_model()
|
||||
|
||||
|
|
Loading…
Reference in New Issue