1 files changed, 67 insertions, 50 deletions

diff --git a/generic/SpatialFullConvolution.c b/generic/SpatialFullConvolution.c
index cb2e340..3b55297 100644
--- a/generic/SpatialFullConvolution.c
+++ b/generic/SpatialFullConvolution.c
@@ -12,69 +12,80 @@ static int nn_(SpatialFullConvolution_updateOutput)(lua_State *L)
   THTensor *bias = luaT_getfieldcheckudata(L, 1, "bias", torch_Tensor);
   THTensor *output = luaT_getfieldcheckudata(L, 1, "output", torch_Tensor);
 
-  luaL_argcheck(L, input->nDimension == 3 || input->nDimension == 4, 2, "3D or 4D(batch mode) tensor expected");
   int dimw = 2;
   int dimh = 1;
+
+  luaL_argcheck(L, input->nDimension == 3 || input->nDimension == 4, 2, "3D or 4D(batch mode) tensor expected");
   if (input->nDimension == 4) {
     dimw++;
     dimh++;
   }
 
-  long nOutputPlane = weight->size[1];
-  long kW           = weight->size[3];
-  long kH           = weight->size[2];
-  long inputWidth   = input->size[dimw];
-  long inputHeight  = input->size[dimh];
-  long outputWidth  = (inputWidth - 1) * dW + kW;
-  long outputHeight = (inputHeight - 1) * dH + kH;
-
-  if (input->nDimension == 3)
-  {
-    THTensor_(resize3d)(output, nOutputPlane, outputHeight, outputWidth);
-    /* add bias */
-    long i;
-    real* bias_data = THTensor_(data)(bias);
-    real* output_data = THTensor_(data)(output);
-#pragma omp parallel for private(i)
-    for (i=0; i<bias->size[0]; i++)
-    {
-      /*THTensor_(select)(outn,output,0,i);*/
-      /*TH_TENSOR_APPLY(real,outn, *outn_data = bias_data[i];);*/
-      real *ptr_output = output_data + i*outputWidth*outputHeight;
-      long j;
-      for(j = 0; j < outputWidth*outputHeight; j++)
-        ptr_output[j] = bias_data[i];
-    }
-
-    /* do convolutions */
-    THTensor *tweight = THTensor_(newTranspose)(weight,0,1);
-    THTensor_(conv2Dmv)(output, 1.0, 1.0, input, tweight, dH, dW, "F", "C");
-    THTensor_(free)(tweight);
-  }
-  else
   {
-    THTensor_(resize4d)(output, input->size[0], nOutputPlane, outputHeight, outputWidth);
-    real* bias_data = THTensor_(data)(bias);
-    real* output_data = THTensor_(data)(output);
-
-    long p;
-#pragma omp parallel for private(p)
-    for (p=0; p<input->size[0]; p++)
+    long nOutputPlane = weight->size[1];
+    long kW           = weight->size[3];
+    long kH           = weight->size[2];
+    long inputWidth   = input->size[dimw];
+    long inputHeight  = input->size[dimh];
+    long outputWidth  = (inputWidth - 1) * dW + kW;
+    long outputHeight = (inputHeight - 1) * dH + kH;
+    
+    if (input->nDimension == 3)
     {
-      /* BIAS */
       long i;
+      real* bias_data;
+      real* output_data;
+
+      THTensor_(resize3d)(output, nOutputPlane, outputHeight, outputWidth);
+      /* add bias */
+      bias_data = THTensor_(data)(bias);
+      output_data = THTensor_(data)(output);
+#pragma omp parallel for private(i)
       for (i=0; i<bias->size[0]; i++)
       {
-        real *ptr_output = output_data + p*nOutputPlane*outputWidth*outputHeight + i*outputWidth*outputHeight;
+        real *ptr_output = output_data + i*outputWidth*outputHeight;
         long j;
         for(j = 0; j < outputWidth*outputHeight; j++)
           ptr_output[j] = bias_data[i];
       }
+      
+      /* do convolutions */
+      {
+        THTensor *tweight = THTensor_(newTranspose)(weight,0,1);
+        THTensor_(conv2Dmv)(output, 1.0, 1.0, input, tweight, dH, dW, "F", "C");
+        THTensor_(free)(tweight);
+      }
+  }
+    else
+    {
+      real* bias_data;
+      real* output_data;
+      long p;
+
+      THTensor_(resize4d)(output, input->size[0], nOutputPlane, outputHeight, outputWidth);
+      bias_data = THTensor_(data)(bias);
+      output_data = THTensor_(data)(output);
+      
+#pragma omp parallel for private(p)
+      for (p=0; p<input->size[0]; p++)
+      {
+        /* BIAS */
+        long i;
+        for (i=0; i<bias->size[0]; i++)
+        {
+          real *ptr_output = output_data + p*nOutputPlane*outputWidth*outputHeight + i*outputWidth*outputHeight;
+          long j;
+          for(j = 0; j < outputWidth*outputHeight; j++)
+            ptr_output[j] = bias_data[i];
+        }
+      }
+      /* do convolutions */
+      {
+        THTensor *tweight = THTensor_(newTranspose)(weight,0,1);
+        THTensor_(conv2Dmm)(output, 1.0, 1.0, input, tweight, dH, dW, "F", "C");
+        THTensor_(free)(tweight);
+      }
     }
-    /* do convolutions */
-    THTensor *tweight = THTensor_(newTranspose)(weight,0,1);
-    THTensor_(conv2Dmm)(output, 1.0, 1.0, input, tweight, dH, dW, "F", "C");
-    THTensor_(free)(tweight);
   }
   return 1;
 }
@@ -120,20 +131,26 @@ static int nn_(SpatialFullConvolution_accGradParameters)(lua_State *L)
   THTensor *gradBias = luaT_getfieldcheckudata(L, 1, "gradBias", torch_Tensor);
   
   long nOutputPlane = weight->size[1];
-  THArgCheck( nOutputPlane == gradOutput->size[input->nDimension == 4 ? 1 : 0], 1, "Number of output features is not equal to nOutputPlane" );
 
   int dimw = 2;
   int dimh = 1;
 
+  real *gradBias_data;
+  real *gradOutput_data;
+  long noutSlice;
+
+  THArgCheck( nOutputPlane == gradOutput->size[input->nDimension == 4 ? 1 : 0], 1, "Number of output features is not equal to nOutputPlane" );
+
+
   if (input->nDimension == 4)
   {
     dimw++;
     dimh++;
   }
   /* gradient to bias */
-  real *gradBias_data = THTensor_(data)(gradBias);
-  real *gradOutput_data = THTensor_(data)(gradOutput);
-  long noutSlice = gradOutput->size[dimh]*gradOutput->size[dimw];
+  gradBias_data = THTensor_(data)(gradBias);
+  gradOutput_data = THTensor_(data)(gradOutput);
+  noutSlice = gradOutput->size[dimh]*gradOutput->size[dimw];
   /*THTensor* gradOutSlice = THTensor_(new)();*/
 
   if (input->nDimension == 3)