From 9d182abadbdec37a3f81ca340bf47144094e41f7 Mon Sep 17 00:00:00 2001
From: thanhvc3 <thanhvc3@vng.com.vn>
Date: Wed, 17 May 2023 13:38:54 +0700
Subject: [PATCH] add grid search

---
 main.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/main.py b/main.py
index 5eccfd9..8986c2d 100644
--- a/main.py
+++ b/main.py
@@ -670,7 +670,7 @@ if __name__ == "__main__":
             sub, rel, obj, label, neg_ent, sub_samp = model.read_batch(
                 batch, 'train')
 
-            dataset = np.stack([sub, rel, neg_ent, np.repeat(model.p.train_strategy, repeats=len(sub))], axis = 1)
+            dataset = np.stack([sub.cpu(), rel.cpu(), neg_ent.cpu(), np.repeat(model.p.train_strategy, repeats=len(sub))], axis = 1)
             search = grid.fit(dataset, label)
             print("BEST SCORE: ", search.best_score_)
             print("BEST PARAMS: ", search.best_params_)