How can we decrease the size of resources for RGATConv?

[songsong0425]

Dear PyG community,

Greetings, always thank you for your great effort for the graph-based deep learning approaches.
I have a simple question about the RGATConv, especially training time and GPU memory usage.

While I am using it for the link prediction task, I feel it is highly dependent on the size of the subgraph generated from the LinkNeighborLoader of the num_neighbors parameter.
Based on the A100 GPU, I frequently encountered the GPU OOM problem when I set the num_neighbors to more than 256.
Here's the code snippet.

parser = argparse.ArgumentParser()
parser.add_argument('--lr', type=float, default=1e-4, help='Learning rate')
parser.add_argument('--num_epochs', type=int, default=200, help='Maximum number of epochs')
parser.add_argument('--batch', type=int, default=4, help='Batch size')
parser.add_argument('--heads', type=int, default=4, help='The number of heads in RGAT')
parser.add_argument('--wd', type=float, default=1e-3, help='Weight decay')
parser.add_argument('--n_neighbor', type=int, default=256, help='The number of sampled neighbors from seed nodes')
args = parser.parse_args()

...
train_loader = LinkNeighborLoader(kf_train_data, edge_label_index=kf_train_data.edge_label_index, edge_label=kf_train_data.edge_label,
                                  batch_size=args.batch, shuffle=True, neg_sampling_ratio=0.0, num_neighbors=[args.n_neighbor,args.n_neighbor//2])
val_loader = LinkNeighborLoader(kf_val_data, edge_label_index=kf_val_data.edge_label_index, edge_label=kf_val_data.edge_label,
                                batch_size=args.batch, shuffle=False, neg_sampling_ratio=0.0, num_neighbors=[args.n_neighbor,args.n_neighbor//2])
test_loader = LinkNeighborLoader(kf_test_data, edge_label_index=kf_test_data.edge_label_index, edge_label=kf_test_data.edge_label,
                                 batch_size=args.batch, shuffle=False, neg_sampling_ratio=0.0, num_neighbors=[args.n_neighbor,args.n_neighbor//2])

model = RGATmodel(args, in_channels=256, hidden_channels=128, out_channels=64, num_relations=KG.num_relations).to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=args.lr, weight_decay=args.wd)
criterion = torch.nn.BCEWithLogitsLoss()

It took almost 80GB of GPU memory and 20 minutes per epoch.
In this case, can we save them by modifying the RGATConv code?
If you have any idea, please let me know.

Thank you for reading this question.