fixed batch TemporalConvolution bug

2 files changed, 14 insertions, 11 deletions

diff --git a/generic/TemporalConvolution.c b/generic/TemporalConvolution.c
index fedb0ce..1805664 100644
--- a/generic/TemporalConvolution.c
+++ b/generic/TemporalConvolution.c
@@ -89,6 +89,7 @@ static int nn_(TemporalConvolution_updateOutput)(lua_State *L)
     {
       THTensor_(select)(outputSample, output, 0, i);
       THTensor_(select)(inputSample, input, 0, i);
+      long nOutputSampleFrame = nOutputFrame;
       
       /* bias first */
       for(k = 0; k < nOutputFrame; k++)
@@ -98,12 +99,12 @@ static int nn_(TemporalConvolution_updateOutput)(lua_State *L)
       }
 
       /* ouch */
-      for(k = 0; nOutputFrame > 0; k++)
+      for(k = 0; nOutputSampleFrame > 0; k++)
       {
         long outputFrameStride = (kW-1)/dW+1;
         long inputFrameStride = outputFrameStride*dW;
         long nFrame = (nInputFrame-k*dW-kW)/inputFrameStride + 1;
-        nOutputFrame -= nFrame;
+        nOutputSampleFrame -= nFrame;
 
         THTensor_(setStorage2d)(inputWindow, inputSample->storage,
                                 inputSample->storageOffset+k*dW*inputSample->size[1],
@@ -198,14 +199,15 @@ static int nn_(TemporalConvolution_updateGradInput)(lua_State *L)
     {
       THTensor_(select)(gradOutputSample, gradOutput, 0, i);
       THTensor_(select)(gradInputSample, gradInput, 0, i);
+      int nOutputSampleFrame = nOutputFrame;
       
       /* ouch */
-      for(k = 0; nOutputFrame > 0; k++)
+      for(k = 0; nOutputSampleFrame > 0; k++)
       {
         long outputFrameStride = (kW-1)/dW+1;
         long inputFrameStride = outputFrameStride*dW;
         long nFrame = (nInputFrame-k*dW-kW)/inputFrameStride + 1;
-        nOutputFrame -= nFrame;
+        nOutputSampleFrame -= nFrame;
 
         THTensor_(setStorage2d)(gradOutputWindow, gradOutputSample->storage,
                                 gradOutputSample->storageOffset + k*gradOutputSample->size[1],
@@ -305,6 +307,7 @@ static int nn_(TemporalConvolution_accGradParameters)(lua_State *L)
     {
       THTensor_(select)(gradOutputSample, gradOutput, 0, i);
       THTensor_(select)(inputSample, input, 0, i);
+      int nOutputSampleFrame = nOutputFrame;
       
       /* bias first */
       for(k = 0; k < nOutputFrame; k++)
@@ -314,15 +317,15 @@ static int nn_(TemporalConvolution_accGradParameters)(lua_State *L)
       }
 
       /* ouch */
-      for(k = 0; nOutputFrame > 0; k++)
+      for(k = 0; nOutputSampleFrame > 0; k++)
       {
         long outputFrameStride = (kW-1)/dW+1;
         long inputFrameStride = outputFrameStride*dW;
         long nFrame = (nInputFrame-k*dW-kW)/inputFrameStride + 1;
-        nOutputFrame -= nFrame;
+        nOutputSampleFrame -= nFrame;
 
         THTensor_(setStorage2d)(inputWindow, inputSample->storage,
-                                input->storageOffset+k*dW*inputSample->size[1],
+                                inputSample->storageOffset+k*dW*inputSample->size[1],
                                 nFrame, inputFrameStride*inputSample->size[1],
                                 kW*inputSample->size[1], 1);
 
diff --git a/test/test.lua b/test/test.lua
index c945d4e..ff4e649 100644
--- a/test/test.lua
+++ b/test/test.lua
@@ -1388,13 +1388,13 @@ function nntest.TemporalConvolution()
    local outputGrad = torch.randn(output:size())
    local inputGrad = module:backward(input, outputGrad):clone()
    
-   local input1D = input:select(1, 1)
+   local input1D = input:select(1, 2)
    local output1D = module:forward(input1D)
-   local outputGrad1D = outputGrad:select(1, 1)
+   local outputGrad1D = outputGrad:select(1, 2)
    local inputGrad1D = module:backward(input1D, outputGrad1D)
    
-   mytester:assertTensorEq(output:select(1,1), output1D, 0.000001, 'error on 2D vs 1D forward)')
-   mytester:assertTensorEq(inputGrad:select(1,1), inputGrad1D, 0.000001, 'error on 2D vs 1D backward)')
+   mytester:assertTensorEq(output:select(1,2), output1D, 0.000001, 'error on 2D vs 1D forward)')
+   mytester:assertTensorEq(inputGrad:select(1,2), inputGrad1D, 0.000001, 'error on 2D vs 1D backward)')
 end
 
 function nntest.TemporalSubSampling()