From 3cfcc4826a2107880e0420ac81291823190ef7fe Mon Sep 17 00:00:00 2001
From: Juhan Bae <pomonam15@gmail.com>
Date: Mon, 8 Jul 2024 17:31:02 -0400
Subject: [PATCH] More debug lines

---
 kronfluence/score/dot_product.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/kronfluence/score/dot_product.py b/kronfluence/score/dot_product.py
index 6421358..99228ba 100644
--- a/kronfluence/score/dot_product.py
+++ b/kronfluence/score/dot_product.py
@@ -103,6 +103,10 @@ def compute_dot_products_with_loader(
             if factor_args.has_shared_parameters:
                 finalize_iteration(model=model, tracked_module_names=tracked_module_names)
 
+            print("Middle")
+            print(torch.cuda.memory_allocated())
+            print(torch.cuda.memory_reserved())
+
             with torch.no_grad():
                 if score_args.compute_per_module_scores:
                     for module in cached_module_lst: