Add debug print

pomonam · Mar 19, 2024 · 0ccec85 · 0ccec85
1 parent 2e84271
commit 0ccec85
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/kronfluence/module/tracked_module.py b/kronfluence/module/tracked_module.py
@@ -679,9 +679,11 @@ def synchronize_preconditioned_gradient(self, num_processes: int) -> None:
         """Stacks preconditioned gradient across multiple devices or nodes in a distributed setting."""
         if dist.is_initialized() and torch.cuda.is_available() and self._preconditioned_gradient_available():
             if isinstance(self._storage[PRECONDITIONED_GRADIENT_NAME], list):
+                print("1")
                 assert len(self._storage[PRECONDITIONED_GRADIENT_NAME]) == 2
                 for i in range(len(self._storage[PRECONDITIONED_GRADIENT_NAME])):
                     size = self._storage[PRECONDITIONED_GRADIENT_NAME][i].size()
+                    print(f"Size: {size}")
                     stacked_matrix = torch.empty(
                         size=(num_processes, size[0], size[1], size[2]),
                         dtype=self._storage[PRECONDITIONED_GRADIENT_NAME][i].dtype,