cleanup: printing messages

moverseai · Jun 7, 2024 · de25301 · de25301
1 parent 1f2f66a
commit de25301
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 3 deletions.
diff --git a/moai/data/datasets/generic/npz.py b/moai/data/datasets/generic/npz.py
@@ -44,7 +44,7 @@ def __init__(
         filename: str = "",
     ):
         self.file = load_npz_file(filename)
-        log.info(f"Loaded an .npz file producing [{list(self.file.keys())}].")
+        log.info(f"Loaded an .npz file producing {list(self.file.keys())}.")
 
     def __len__(self) -> int:
         return len(self.file[toolz.first(self.file)])
@@ -63,7 +63,7 @@ def __init__(
     ):
         self.file = load_npz_file(filename)
         self.length = length
-        log.info(f"Loaded an .npz file producing [{list(self.file.keys())}].")
+        log.info(f"Loaded an .npz file producing {list(self.file.keys())}.")
 
     def __len__(self) -> int:
         return self.length

diff --git a/moai/parameters/initialization/schemes/zero_flow_params.py b/moai/parameters/initialization/schemes/zero_flow_params.py
@@ -18,13 +18,16 @@ def __init__(
         self.keys = keys
 
     def __call__(self, module: torch.nn.Module) -> None:
+        zeroed_keys = []
         for key in self.keys:
             try:
                 m = get_parameter(module.named_flows, key)
                 if m is not None:
-                    log.info(f"Zeroing out parameter: [cyan italic]{key}[/].")
                     with torch.no_grad():  # TODO: remove this and add in root apply call
                         m.zero_()
                         m.grad = None
+                zeroed_keys.append(key)
             except:
                 break
+        all_zeroed_keys = ",".join(zeroed_keys)
+        log.info(f"Zeroing out parameters: [cyan italic]\[{all_zeroed_keys}][/].")