Update run script with torchinfo

1 files changed, 6 insertions, 4 deletions

diff --git a/training/run.py b/training/run.py
index ddae2b9..b854376 100644
--- a/training/run.py
+++ b/training/run.py
@@ -13,7 +13,7 @@ from pytorch_lightning import (
 )
 from pytorch_lightning.loggers import LightningLoggerBase
 from torch import nn
-from torchsummary import summary
+from torchinfo import summary
 from text_recognizer.data.base_mapping import AbstractMapping
 
 import utils
@@ -38,9 +38,6 @@ 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)
 
@@ -69,6 +66,11 @@ def run(config: DictConfig) -> Optional[float]:
     utils.log_hyperparameters(config=config, model=model, trainer=trainer)
     utils.save_config(config)
 
+    if config.get("summary"):
+        summary(
+            network, list(map(lambda x: list(x), config.summary)), depth=1, device="cpu"
+        )
+
     if config.debug:
         log.info("Fast development run...")
         trainer.fit(model, datamodule=datamodule)