training time

pyg-team · Oct 1, 2024 · bd25931 · bd25931
1 parent 5e29d9c
commit bd25931
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/examples/llm/ogbg_code2.py b/examples/llm/ogbg_code2.py
@@ -44,8 +44,8 @@ def inference_step_ogbg(model, batch, **kwargs):
 )
 
 # TODO, try more epochs
-train(num_epochs=1, hidden_channels=None, num_gnn_layers=None, batch_size=12,
-      eval_batch_size=24, lr=1e-5, checkpointing=False, model=CodeRetriever,
+train(num_epochs=1, hidden_channels=None, num_gnn_layers=None, batch_size=8,
+      eval_batch_size=16, lr=1e-5, checkpointing=False, model=CodeRetriever,
       dataset=OGBG_Code2, get_loss=get_loss_ogbg,
       inference_step=inference_step_ogbg, model_save_name="code_retriever")
 torch.cuda.empty_cache()