Reduce size

pomonam · Jul 8, 2024 · 386b335 · 386b335
1 parent 3cfcc48
commit 386b335
Show file tree

Hide file tree

Showing 2 changed files with 1 addition and 2 deletions.
diff --git a/kronfluence/module/linear.py b/kronfluence/module/linear.py
@@ -103,7 +103,7 @@ def compute_pairwise_score(
                 minimize = "size"
             else:
                 expr = "qio,b...i,b...o->qb"
-                minimize = "flops"
+                minimize = "size"
             self.einsum_expression = contract_expression(
                 expr,
                 preconditioned_gradient.shape,

diff --git a/kronfluence/score/dot_product.py b/kronfluence/score/dot_product.py
@@ -128,7 +128,6 @@ def compute_dot_products_with_loader(
                             raise
                     pairwise_scores = pairwise_scores.cpu()
                     score_chunks[ALL_MODULE_NAME].append(pairwise_scores)
-                    del pairwise_scores
                     accumulate_iterations(model=model, tracked_module_names=tracked_module_names)
 
             if state.use_distributed and total_steps % DISTRIBUTED_SYNC_INTERVAL == 0: