Improving the JKNET example. (#6596)

09a1a2f8 · Andrei Ivanov · GitHub · f6345635 · 09a1a2f8
Unverified Commit 09a1a2f8 authored Nov 21, 2023 by Andrei Ivanov Committed by GitHub Nov 22, 2023
Hide whitespace changes
Inline Side-by-side

Showing with 3 additions and 0 deletions

examples/pytorch/jknet/model.py examples/pytorch/jknet/model.py +3 -0

No files found.
--- a/examples/pytorch/jknet/model.py
+++ b/examples/pytorch/jknet/model.py
@@ -41,6 +41,9 @@ class JKNet(nn.Module):
            feats = self.dropout(layer(g, feats))
            feat_lst.append(feats)

+        if self.mode == "lstm":
+            self.jump.lstm.flatten_parameters()
+
        g.ndata["h"] = self.jump(feat_lst)
        g.update_all(fn.copy_u("h", "m"), fn.sum("m", "h"))