mLSTM cpu fix

435133f1 · Raul Puri · GitHub · 83acda92 · 435133f1
Unverified Commit 435133f1 authored May 15, 2018 by Raul Puri Committed by GitHub May 15, 2018
Hide whitespace changes
Inline Side-by-side

Showing with 1 addition and 1 deletion

apex/RNN/cells.py apex/RNN/cells.py +1 -1

No files found.
--- a/apex/RNN/cells.py
+++ b/apex/RNN/cells.py
@@ -68,7 +68,7 @@ def mLSTMCell(input, hidden, w_ih, w_hh, w_mih, w_mhh, b_ih=None, b_hh=None):
    hx, cx = hidden
    
    m = F.linear(input, w_mih) * F.linear(hidden[0], w_mhh)
-    igates = F.linear(input, w_ih, b_ih) + F.linear(m, w_hh, b_hh)
+    gates = F.linear(input, w_ih, b_ih) + F.linear(m, w_hh, b_hh)

    ingate, forgetgate, cellgate, outgate = gates.chunk(4, 1)