mirror of
https://github.com/wassname/pytorch-lightning.git
synced 2026-10-30 13:40:46 +08:00
added sample input for summary
This commit is contained in:
1 parent
f3b0cbf998
commit
83ccd21bec
1 file changed
+2
@@ -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()
|
||||
|
||||
|
||||
Reference in new issue
Block a user