fix Conv3d non-contiguous weight bug

1 files changed, 13 insertions, 16 deletions

diff --git a/lib/THNN/generic/VolumetricConvolutionMM.c b/lib/THNN/generic/VolumetricConvolutionMM.c
index 00a121d..1d7928b 100644
--- a/lib/THNN/generic/VolumetricConvolutionMM.c
+++ b/lib/THNN/generic/VolumetricConvolutionMM.c
@@ -82,16 +82,18 @@ static void inline THNN_(VolumetricConvolutionMM_shapeCheck)(
   }
 }
 
-static int THNN_(view_weight)(THTensor **_weight)
+static THTensor* THNN_(view_weight)(THTensor *weight)
 {
-  THTensor *weight = *_weight;
+  weight = THTensor_(newContiguous)(weight);
   if (weight->nDimension == 5) {
     long s1 = weight->size[0];
     long s2 = weight->size[1] * weight->size[2] * weight->size[3] * weight->size[4];
-    *_weight = THTensor_(newWithStorage2d)(weight->storage, weight->storageOffset, s1, -1, s2, -1);
-    return 1;
+    THTensor *old_weight = weight;
+    weight = THTensor_(newWithStorage2d)(weight->storage, weight->storageOffset,
+					 s1, -1, s2, -1);
+    THTensor_(free)(old_weight);
   }
-  return 0;
+  return weight;
 }
 
 /* note: due to write issues, this one cannot be parallelized as well as unfolded_copy */
@@ -341,7 +343,6 @@ void THNN_(VolumetricConvolutionMM_updateOutput)(
   int dimt = 1;
   int dimh = 2;
   int dimw = 3;
-  int freeWeight = 0;
 
   long nInputPlane;
   long inputDepth;
@@ -374,7 +375,7 @@ void THNN_(VolumetricConvolutionMM_updateOutput)(
   outputHeight = (inputHeight + 2*pH - kH) / dH + 1;
   outputWidth  = (inputWidth + 2*pW - kW) / dW + 1;
 
-  freeWeight = THNN_(view_weight)(&weight);
+  weight = THNN_(view_weight)(weight);
 
   if (input->nDimension == 4)
   {
@@ -421,8 +422,7 @@ void THNN_(VolumetricConvolutionMM_updateOutput)(
   }
 
   THTensor_(free)(input);
-  if (freeWeight)
-    THTensor_(free)(weight);
+  THTensor_(free)(weight);
 }
 
 static void THNN_(VolumetricConvolutionMM_updateGradInput_frame)(
@@ -487,7 +487,7 @@ void THNN_(VolumetricConvolutionMM_updateGradInput)(
   input = THTensor_(newContiguous)(input);
   gradOutput = THTensor_(newContiguous)(gradOutput);
 
-  int freeWeight = THNN_(view_weight)(&weight);
+  weight = THNN_(view_weight)(weight);
 
   THTensor_(resizeAs)(gradInput, input);
   THTensor_(resizeAs)(fgradInput, finput);
@@ -535,8 +535,7 @@ void THNN_(VolumetricConvolutionMM_updateGradInput)(
   THTensor_(free)(tweight);
   THTensor_(free)(input);
   THTensor_(free)(gradOutput);
-  if (freeWeight)
-    THTensor_(free)(weight);
+  THTensor_(free)(weight);
 }
 
 static void THNN_(VolumetricConvolutionMM_accGradParameters_frame)(
@@ -587,7 +586,6 @@ void THNN_(VolumetricConvolutionMM_accGradParameters)(
           accreal scale_)
 {
   real scale = TH_CONVERT_ACCREAL_TO_REAL(scale_);
-  int freeWeight;
   int nOutputPlane = (int)gradWeight->size[0];
 
   THNN_(VolumetricConvolutionMM_shapeCheck)(
@@ -596,7 +594,7 @@ void THNN_(VolumetricConvolutionMM_accGradParameters)(
   input = THTensor_(newContiguous)(input);
   gradOutput = THTensor_(newContiguous)(gradOutput);
 
-  freeWeight = THNN_(view_weight)(&gradWeight);
+  gradWeight = THNN_(view_weight)(gradWeight);
 
   if (input->nDimension == 4)   // non-batch mode
   {
@@ -621,8 +619,7 @@ void THNN_(VolumetricConvolutionMM_accGradParameters)(
 
   THTensor_(free)(input);
   THTensor_(free)(gradOutput);
-  if (freeWeight)
-    THTensor_(free)(gradWeight);
+  THTensor_(free)(gradWeight);
 }
 
 #endif