Issue with multi GPUs training

[will-deeplearn]

I am trying to train on the tpugraphs dataset, and I encounter an OOM error during the validation step.

Start from epoch 0 train: {'epoch': 0, 'time_epoch': 17.86623, 'eta': 357306.6434, 'eta_hours': 99.25185, 'loss': 49.00541394, 'lr': 0.0001, 'params': 546197, 'time_iter': 1.98514, 'opa': 0.52127, 'spearmanr': 0.06584} ...computing epoch stats took: 0.04s val: {'epoch': 0, 'time_epoch': 1.97774, 'loss': 49.72499466, 'lr': 0, 'params': 546197, 'time_iter': 1.97774, 'opa': 0.45795, 'spearmanr': -0.12068} ...computing epoch stats took: 0.01s

I am unsure how to resolve this issue. When I set the batch size (bs) to 1, I get another error.

Traceback (most recent call last): File "GST/main_tpugraphs.py", line 173, in <module> train_dict[cfg.train.mode](loggers, loaders, model, optimizer, File "/GST/graphgps/train/custom_tpu_train.py", line 337, in custom_train eval_epoch(loggers[i], loaders[i], model, File "miniconda3/envs/graphgps/lib/python3.10/site-packages/torch/autograd/grad_mode.py", line 27, in decorate_context return func(*args, **kwargs) File "GST/graphgps/train/custom_tpu_train.py", line 267, in eval_epoch logger.update_stats(true=_true, File "GST/graphgps/logger.py", line 232, in update_stats assert true.shape[0] == pred.shape[0] AssertionError

I am considering training across multiple GPUs to solve the OOM issue. I have tried setting cfg.devices to 2 and using DataParallel, but I cannot achieve multi-GPU training. Can you help address this issue? Additionally, I found a bug in custom_tpu_train.py. We should iterate over model.model.model.children() instead of model.model.children().