Merge pull request #207 from nhynes/batchfirst-untranspose

Un-untranspose RNN outputs after doing grad updates

1 files changed, 7 insertions, 0 deletions

diff --git a/RNN.lua b/RNN.lua
index daf3b19..23b1d46 100644
--- a/RNN.lua
+++ b/RNN.lua
@@ -437,6 +437,7 @@ function RNN:updateGradInput(input, gradOutput)
             self.reserve:data(), self.reserve:size(1) * 4) -- sizeof(float)
     if (self.batchFirst) then
         self.gradInput = self.gradInput:transpose(1, 2)
+        self.output = self.output:transpose(1, 2)
     end
    return self.gradInput
 end
@@ -445,6 +446,7 @@ function RNN:accGradParameters(input, gradOutput, scale)
     if (self.batchFirst) then
         input = input:transpose(1, 2)
         gradOutput = gradOutput:transpose(1, 2)
+        self.output = self.output:transpose(1, 2)
     end
    scale = scale or 1
    if scale == 0 then return end
@@ -502,6 +504,11 @@ function RNN:accGradParameters(input, gradOutput, scale)
                self.dw:data(),
                scaleTensor:data())
    end
+
+    if (self.batchFirst) then
+        gradOutput = gradOutput:transpose(1, 2)
+        self.output = self.output:transpose(1, 2)
+    end
 end
 
 function RNN:clearDesc()