Better styling

Lightning-AI · Sep 17, 2021 · 6dd81bc · 6dd81bc
1 parent a4d6abb
commit 6dd81bc
Showing 1 changed file with 6 additions and 7 deletions.
diff --git a/pytorch_lightning/callbacks/rich_model_summary.py b/pytorch_lightning/callbacks/rich_model_summary.py
@@ -75,18 +75,17 @@ def summarize(
 
         console = Console()
 
-        table = Table(title="Model Summary")
-
-        table.add_column(" ")
-        table.add_column("Name", justify="left", style="cyan", no_wrap=True)
-        table.add_column("Type", style="magenta")
-        table.add_column("Params", justify="right", style="green")
+        table = Table(header_style="bold magenta")
+        table.add_column(" ", style="dim")
+        table.add_column("Name", justify="left", no_wrap=True)
+        table.add_column("Type")
+        table.add_column("Params", justify="right")
 
         column_names = list(zip(*summary_data))[0]
 
         for column_name in ["In sizes", "Out sizes"]:
             if column_name in column_names:
-                table.add_column(column_name, justify="right", style="green")
+                table.add_column(column_name, justify="right", style="white")
 
         rows = list(zip(*(arr[1] for arr in summary_data)))
         for row in rows: