diff --git a/main.py b/main.py
index 839e9c5..a421f7b 100644
--- a/main.py
+++ b/main.py
@@ -653,10 +653,10 @@ if __name__ == "__main__":
             'optimizer__lr': [0.0001, 0.0003, 0.001],
             'optimizer__weight_decay': [1e-4, 1e-5, 1e-6],
             'module__hid_drop': [0.2, 0.5, 0.7],
-            'module__embed_dim': [300, 400, 500],
+            # 'module__embed_dim': [300, 400, 500],
         }
 
-        grid = GridSearchCV(estimator=estimator, param_grid=paramsGrid, n_jobs=-1, cv=3, scoring=torch.nn.BCELoss)
+        grid = GridSearchCV(estimator=estimator, param_grid=paramsGrid, n_jobs=-1, cv=2, scoring=torch.nn.BCELoss)
         data = np.array(model.triples['train'])
         data = data[np.random.choice(np.arange(len(data)), size=int(len(data) * 0.15), replace=False)]
         dataloader = iter(DataLoader(