Updated gradOutput assertion

1 files changed, 2 insertions, 8 deletions

diff --git a/RNN.lua b/RNN.lua
index 1b945fc..af1272d 100644
--- a/RNN.lua
+++ b/RNN.lua
@@ -365,15 +365,9 @@ function RNN:updateGradInput(input, gradOutput)
    assert(input:size(1) == self.seqLength, 'input has incorrect sequence length!')
    assert(input:size(2) == self.miniBatch, 'input has incorrect minibatch size!')
    assert(input:size(3) == self.inputSize, 'input has incorrect size!')
-
-    if (self.batchFirst) then
-        assert(gradOutput:isSameSizeAs(self.output:transpose(1, 2)), 'gradOutput has incorrect size!')
-
-    else
-        assert(gradOutput:isSameSizeAs(self.output), 'gradOutput has incorrect size!')
-    end
    assert(self.train, 'updateGradInput can only be called when training!')
-
+   local expectedSize = torch.LongStorage {self.seqLength, self.miniBatch, self.hiddenSize * self.numDirections}
+   assert(gradOutput:isSize(expectedSize), 'gradOutput has incorrect size!')
    local x, dy = self:makeContiguous(input, gradOutput)
    local y = self.output
    local w = self.weight