Renamed --print-density argument to --print-model

Rob Cornish · Rob Cornish · commit a8c7b4cc42cc · 2019-11-06T21:12:18.000Z
diff --git a/lgf/experiment.py b/lgf/experiment.py
@@ -31,7 +31,7 @@ def train(config):
         trainer.train()
 
 
-def print_density(config):
+def print_model(config):
     _, density, _, _ = setup_experiment({**config, "enable_logging": False})
     print(density)
     print(f"Number of parameters: {num_params(density):,}")
diff --git a/main.py b/main.py
@@ -18,7 +18,7 @@
 ], required=True)
 parser.add_argument("--baseline", action="store_true", help="Run baseline flow instead of LGF")
 parser.add_argument("--seed", type=int, help="Random seed to use. Defaults to using current time.")
-parser.add_argument("--print-density", action="store_true", help="Print the Pytorch Density and exit")
+parser.add_argument("--print-model", action="store_true", help="Print the model and exit")
 parser.add_argument("--print-schema", action="store_true", help="Print the model schema and exit")
 parser.add_argument("--print-config", action="store_true", help="Print the full config and exit")
 parser.add_argument("--nochkpt", action="store_true", help="Disable checkpointing")
@@ -77,9 +77,9 @@ def parse_config_arg(key_value):
     print(json.dumps(config, indent=4))
     should_train = False
 
-if args.print_density:
-    from lgf.experiment import print_density
-    print_density({**config, "write_to_disk": False})
+if args.print_model:
+    from lgf.experiment import print_model
+    print_model({**config, "write_to_disk": False})
     should_train = False
 
 if args.print_schema:
