Skip to content

GitLab

"vscode:/vscode.git/clone" did not exist on "3e06bc6f2ddd56cd713229a7f0404ff56cabda47"
Unverified Commit 09a1a2f8 authored by Andrei Ivanov's avatar Andrei Ivanov Committed by GitHub
Browse files

Improving the JKNET example. (#6596)

parent f6345635
Show whitespace changes
Inline Side-by-side
Showing with 3 additions and 0 deletions
examples/pytorch/jknet/model.py
View file @ 09a1a2f8
......@@ -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"))
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment