Fixed bug in Sqrt/Square

2 files changed, 17 insertions, 55 deletions

diff --git a/generic/Sqrt.c b/generic/Sqrt.c
index 952c260..d40918b 100644
--- a/generic/Sqrt.c
+++ b/generic/Sqrt.c
@@ -19,19 +19,10 @@ static int nn_(Sqrt_updateOutput)(lua_State *L)
   {
     real* output_data = THTensor_(data)(output);
     real* input_data  = THTensor_(data)(input);
-    long k;
-
-#pragma omp parallel for private(k)
-    for (k = 0; k < input->size[0]; k++)
-    {
-      real* ptr_output = output_data + k*input->stride[0];
-      real* ptr_input  = input_data  + k*input->stride[0];
-      long i;
-      for (i = 0; i < input->stride[0]; i++)
-      {
-        ptr_output[i] = sqrt(ptr_input[i] + bias);
-      }
-    }
+    long i;
+#pragma omp parallel for private(i)
+    for(i = 0; i < THTensor_(nElement)(input); i++)
+      output_data[i] = sqrt(input_data[i] + bias);
   }
   return 1;
 }
@@ -58,20 +49,10 @@ static int nn_(Sqrt_updateGradInput)(lua_State *L)
     real* gradOutput_data = THTensor_(data)(gradOutput);
     real* gradInput_data  = THTensor_(data)(gradInput);
     real* output_data     = THTensor_(data)(output);
-    long k;
-
-#pragma omp parallel for private(k)
-    for (k = 0; k < output->size[0]; k++)
-    {
-      real* ptr_gradOutput = gradOutput_data + k*output->stride[0];
-      real* ptr_gradInput  = gradInput_data  + k*output->stride[0];
-      real* ptr_output     = output_data     + k*output->stride[0];
-      long i;
-      for (i = 0; i < output->stride[0]; i++)
-      {
-        ptr_gradInput[i] = 0.5 * (ptr_gradOutput[i] / ptr_output[i]);
-      }
-    }
+    long i;
+#pragma omp parallel for private(i)
+    for(i = 0; i < THTensor_(nElement)(output); i++)
+      gradInput_data[i] = 0.5 * (gradOutput_data[i] / output_data[i]);
   }
   return 1;
 }
diff --git a/generic/Square.c b/generic/Square.c
index 97baee3..ba0c162 100644
--- a/generic/Square.c
+++ b/generic/Square.c
@@ -18,19 +18,10 @@ static int nn_(Square_updateOutput)(lua_State *L)
   {
     real* output_data = THTensor_(data)(output);
     real* input_data  = THTensor_(data)(input);
-    long k;
-
-#pragma omp parallel for private(k)
-    for (k = 0; k < input->size[0]; k++)
-    {
-      real* ptr_output = output_data + k*input->stride[0];
-      real* ptr_input  = input_data  + k*input->stride[0];
-      long i;
-      for (i = 0; i < input->stride[0]; i++)
-      {
-        ptr_output[i] = ptr_input[i]*ptr_input[i];
-      }
-    }
+    long i;
+#pragma omp parallel for private(i)
+    for(i = 0; i < THTensor_(nElement)(input); i++)
+      output_data[i] = input_data[i]*input_data[i];
   }
   return 1;
 }
@@ -55,21 +46,11 @@ static int nn_(Square_updateGradInput)(lua_State *L)
   {
     real* gradOutput_data = THTensor_(data)(gradOutput);
     real* gradInput_data  = THTensor_(data)(gradInput);
-    real* input_data     = THTensor_(data)(input);
-    long k;
-
-#pragma omp parallel for private(k)
-    for (k = 0; k < input->size[0]; k++)
-    {
-      real* ptr_gradOutput = gradOutput_data + k*input->stride[0];
-      real* ptr_gradInput  = gradInput_data  + k*input->stride[0];
-      real* ptr_input     = input_data     + k*input->stride[0];
-      long i;
-      for (i = 0; i < input->stride[0]; i++)
-      {
-        ptr_gradInput[i] = 2.0 * ptr_gradOutput[i] * ptr_input[i];
-      }
-    }
+    real* input_data  = THTensor_(data)(input);
+    long i;
+#pragma omp parallel for private(i)
+    for(i = 0; i < THTensor_(nElement)(gradInput); i++)
+      gradInput_data[i] = 2.0 * gradOutput_data[i] * input_data[i];
   }
   return 1;
 }