From 786197ab1cc6a6d8ae296009b47cfa69efc0e197 Mon Sep 17 00:00:00 2001 From: Muhammad Dawood Date: Mon, 5 Jun 2023 11:01:42 +0100 Subject: [PATCH] Update gnn.py GPU Out of memory issue problem. --- workspace/model/gnn.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/workspace/model/gnn.py b/workspace/model/gnn.py index fd81a2b..0bd20db 100644 --- a/workspace/model/gnn.py +++ b/workspace/model/gnn.py @@ -118,8 +118,9 @@ def forward(self, data): import torch.nn.functional as F for layer in range(self.no_layers): if layer == 0: + # Uncomment this line and use smaller size for node level features if you are + # facing GPU memory problem. #x = self.first_h(x) - # import pdb; pdb.set_trace() z = self.linears[layer](x) Z += z dout = F.dropout(pooling(z, batch), @@ -455,4 +456,4 @@ def train(self, train_loader, max_epochs=100, optimizer=torch.optim.Adam, schedu test_pr = test_pr_at_best_val_acc Q.reverse() - return Q, train_loss, train_acc, val_loss, np.round(val_acc, 2), test_loss, np.round(test_acc, 2), val_pr, test_pr \ No newline at end of file + return Q, train_loss, train_acc, val_loss, np.round(val_acc, 2), test_loss, np.round(test_acc, 2), val_pr, test_pr