19 lines
465 B
Bash
19 lines
465 B
Bash
CUDA_VISIBLE_DEVICES=0 python train_tom.py \
|
|
--model_type graphbcrnn \
|
|
--types single_object preference instrumental_action \
|
|
--data_path /datasets/external/bib_train/graphs/all_tasks/ \
|
|
--seed 7 \
|
|
--batch_size 32 \
|
|
--max_epochs 35 \
|
|
--gpus 1 \
|
|
--auto_select_gpus True \
|
|
--num_workers 2 \
|
|
--stochastic_weight_avg True \
|
|
--lr 5e-4 \
|
|
--check_val_every_n_epoch 1 \
|
|
--track_grad_norm 2 \
|
|
--gradient_clip_val 10 \
|
|
--gnn_type RSAGEv4 \
|
|
--state_dim 96 \
|
|
--aggregation sum
|