From 18b410fd7e4003116b46d9e5847a0cb70e913c42 Mon Sep 17 00:00:00 2001
From: Hi-archer <285484026@QQ.COM>
Date: Thu, 5 Dec 2024 10:34:53 +0800
Subject: [PATCH] Code Bug Fix

---
 easyeditor/models/rome/rome_main.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/easyeditor/models/rome/rome_main.py b/easyeditor/models/rome/rome_main.py
index fa63b15..f69bb3e 100644
--- a/easyeditor/models/rome/rome_main.py
+++ b/easyeditor/models/rome/rome_main.py
@@ -31,6 +31,7 @@ def set_min_to_zero(pos_matrix, matrix, possibility):
     indices_to_zero = np.argpartition(flattened_pos_matrix, num_zeros)[:num_zeros]
     flattened_matrix[indices_to_zero] = 0
     modified_matrix = flattened_matrix.reshape(matrix.shape)
+    return modified_matrix
 
 def apply_rome_to_model(
     model: AutoModelForCausalLM,