New Model running

2020-06-23 14:37:34 +02:00
parent a19bd9cafd
commit 1033b26195
12 changed files with 173 additions and 112 deletions
--- a/main.py
+++ b/main.py
@ -46,6 +46,7 @@ def run_lightning_loop(config_obj):
        # Init
        model: LightningBaseModule = config_obj.model_class(config_obj.model_paramters)
        model.init_weights(torch.nn.init.xavier_normal_)
+        model.save_to_disk(logger.log_dir)

        # Trainer
        # =============================================================================