Skip to content

Commit

Permalink
[brief] Adds a way to disable the banner.
Browse files Browse the repository at this point in the history
[detailed]
- Mostly for aesthetic reasons.
  • Loading branch information
marovira committed May 7, 2024
1 parent 3ba57cd commit ced25a9
Showing 1 changed file with 5 additions and 1 deletion.
6 changes: 5 additions & 1 deletion src/helios/trainer.py
Original file line number Diff line number Diff line change
Expand Up @@ -198,6 +198,7 @@ class Trainer:
chkpt_root (pathlib.Path): root folder in which checkpoints will be placed.
log_path (pathlib.Path): root folder in which logs will be saved.
run_path (pathlib.Path): root folder in which Tensorboard runs will be saved.
print_banner (bool): if True, the Helios banner with system info will be printed.
"""

def __init__(
Expand All @@ -223,6 +224,7 @@ def __init__(
run_path: pathlib.Path | None = None,
src_root: pathlib.Path | None = None,
import_prefix: str = "",
print_banner: bool = True,
):
"""Create the trainer."""
self._model: hlm.Model | None = None
Expand Down Expand Up @@ -263,6 +265,7 @@ def __init__(
self._import_prefix = import_prefix

self._run_name = run_name
self._print_banner = print_banner

self._validate_flags()
self._setup_device_flags(use_cpu)
Expand Down Expand Up @@ -545,7 +548,8 @@ def _print_header(
"""Print the Helios header with system info to the logs."""
root_logger = logging.get_root_logger()

dist.global_print(core.get_env_info_str())
if self._print_banner:
dist.global_print(core.get_env_info_str())

if for_training:
if chkpt_path is not None:
Expand Down

0 comments on commit ced25a9

Please sign in to comment.