mirror of
https://github.com/wassname/pytorch-lightning.git
synced 2026-09-09 11:32:07 +08:00
added sample input for summary
This commit is contained in:
@@ -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