   int outputSize = luaT_getfieldcheckint(L, 1, "outputSize");
   int inputWindowSize = luaT_getfieldcheckint(L, 1, "inputWindowSize");
   int outputWindowSize = luaT_getfieldcheckint(L, 1, "outputWindowSize");
+  int windowStride = luaT_getfieldcheckint(L, 1, "windowStride");
   int batchSize = luaT_getfieldcheckint(L, 1, "batchSize");
-  int softmax = luaT_getfieldcheckboolean(L, 2, "softmax");
+  int softmax = luaT_getfieldcheckboolean(L, 1, "softmax");
   
+  THCudaTensor *inputIndiceCuda = (THCudaTensor*)luaT_getfieldcheckudata(L, 1, "inputIndiceCuda", "torch.CudaTensor");
+  THLongTensor *inputIndice = (THLongTensor*)luaT_getfieldcheckudata(L, 1, "inputIndice", "torch.LongTensor");
   THCudaTensor *outputIndiceCuda = (THCudaTensor*)luaT_getfieldcheckudata(L, 1, "outputIndiceCuda", "torch.CudaTensor");
   THLongTensor *outputIndice = (THLongTensor*)luaT_getfieldcheckudata(L, 1, "outputIndice", "torch.LongTensor");
   THCudaTensor *output = (THCudaTensor*)luaT_getfieldcheckudata(L, 1, "output", "torch.CudaTensor");
+  
+  luaL_argcheck(L, input->nDimension == 2, 2, "2D(batch mode) tensor expected");
+  luaL_argcheck(L, input->size[1] == inputSize, 2, "invalid input size"); 
+  
+  THCudaTensor_resize1d(output, batchSize, outputWindowSize);
+  THCudaTensor_resize1d(inputIndiceCuda, batchSize);
+  THCudaTensor_resize1d(outputIndiceCuda, batchSize);
     
   /* call cudakernel */
   dim3 blocks(batchSize); // each cuda-block is an example
   dim3 threads(WINDOWGATE_THREADS);
   cunnx_WindowGate_updateOutput_kernel<<<blocks,threads>>>(
-    THCudaTensor_data(output), THCudaTensor_data(outputIndiceCuda),
-    (const float*)THCudaTensor_data(input), 
-    inputWindowSize, outputWindowSize, inputSize, outputSize,
-    softmax
+    THCudaTensor_data(output), THCudaTensor_data(inputIndiceCuda),
+    THCudaTensor_data(outputIndiceCuda),
+    (const float*)THCudaTensor_data(input), inputSize, outputSize,
+    inputWindowSize, outputWindowSize, windowStride, softmax
   );
   
+  THLongTensor_copyCuda(inputIndice, inputIndiceCuda);
+  THLongTensor_copyCuda(outputIndice, outputIndiceCuda);
 
-  return 1;
+  return 0;
 }
 
 

      buffer[tx] += buffer[tx+stride];
  }
  
  __syncthreads();
  float centroid = buffer[0];
  
  if (!softmax) 
    centroid /=inputSize;
  
  // make centroids a number between 0 and 1
  centroid /=inputSize;
  
  self._inputIndice:mul(self.centroid, input:size(2)):add(-self.inputWindowSize*0.5)
  self.inputIndice:copy(self._inputIndice)
  self._outputIndice:mul(self.centroid, self.outputSize):add(-self.outputWindowSize*0.5)
  self.outputIndice:copy(self._outputIndice)
}

  

#define WINDOWGATE_THREADS 128

__global__ void cunnx_WindowGate_updateOutput_kernel(
  float *output, float *inputIndiceCuda, float *outputIndiceCuda,
  const float* input, int inputSize, int outputSize, int inputWindowSize, 
  int outputWindowSize, int windowStride, int softmax)
{
  __shared__ float buffer[WINDOWGATE_THREADS+1];
  unsigned int tx = threadIdx.x;
  unsigned int k = blockIdx.x;
  const float *input_k = input + inputSize*k;
  float *output_k = output + outputWindowSize*k;
  int inputIdx;
  int outputIdx;
  
  // get coordinate of centoid
  buffer[tx] = 0;
