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