add grid search

2023-05-17 13:48:49 +07:00
parent bb6d06a903
commit 835213ab1b
1 changed files with 1 additions and 1 deletions
--- a/main.py
+++ b/main.py
@@ -656,7 +656,7 @@ if __name__ == "__main__":
            'module__embed_dim': [300, 400, 500],
        }

-        grid = GridSearchCV(estimator=estimator, param_grid=paramsGrid, n_jobs=-1, cv=1, scoring=torch.nn.BCELoss)
+        grid = GridSearchCV(estimator=estimator, param_grid=paramsGrid, n_jobs=-1, cv=3, scoring=torch.nn.BCELoss)
        data = np.array(model.triples['train'])
        data = data[np.random.choice(np.arange(len(data)), size=int(len(data) * 0.2), replace=False)]
        dataloader = iter(DataLoader(