From e5a343b0c5a5c3dd87762fc785aa2d1a87b4c2e5 Mon Sep 17 00:00:00 2001 From: thanhvc3 Date: Wed, 17 May 2023 13:37:53 +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 2626de1..5eccfd9 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)], axis = 1) + dataset = np.stack([sub, rel, neg_ent, 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_)