Fix VQVAE into en/decoder, bug in wandb artifact code uploading

1 files changed, 4 insertions, 0 deletions

diff --git a/training/run.py b/training/run.py
index 13a6a82..a2529b0 100644
--- a/training/run.py
+++ b/training/run.py
@@ -13,6 +13,7 @@ from pytorch_lightning import (
 )
 from pytorch_lightning.loggers import LightningLoggerBase
 from torch import nn
+from torchsummary import summary
 from text_recognizer.data.base_mapping import AbstractMapping
 
 import utils
@@ -37,6 +38,9 @@ def run(config: DictConfig) -> Optional[float]:
     log.info(f"Instantiating network <{config.network._target_}>")
     network: nn.Module = hydra.utils.instantiate(config.network)
 
+    if config.summary:
+        summary(network, tuple(config.summary), device="cpu")
+
     log.info(f"Instantiating criterion <{config.criterion._target_}>")
     loss_fn: Type[nn.Module] = hydra.utils.instantiate(config.criterion)