         THCudaTensor *output_n = THCudaTensor_new();
 
         // For each elt in batch, do:
-        for (int elt = 0; elt < batchSize; elt ++) {
+        for (int elt = 0; elt < batchSize; elt += stepBatchSize) {
             // Extract columns:
             imt2col(
                 THCudaTensor_data(input) + elt * inputHeight * inputWidth * nInputPlane,
-                nInputPlane, inputHeight, inputWidth, kW, padding, dW, 
+                nInputPlane, inputHeight, inputWidth, kW, padding, dW, stepBatchSize,
                 THCudaTensor_data(columns)
             );
 
             // Matrix mulitply per output:
-            THCudaTensor_select(output_n, output, 0, elt);
+            THCudaTensor_narrow(output_n, output, 0, elt, stepBatchSize);
 
             // M,N,K are dims of matrix A and B
             // (see http://docs.nvidia.com/cuda/cublas/#cublas-lt-t-gt-gemm)
-            long m = output_n->size[2];
-            long n = output_n->size[0] * output_n->size[1];
+            long m = weight->size[0];
+            long n = columns->size[1];
             long k = weight->size[1];
 
             // Do GEMM_BHWD (note: this is a bit confusing because gemm assumes column-major matrices)

        const int height, const int width, const int ksize, const int pad,
        const int stride, const int channels,
        const int height_col, const int width_col,
        float* data_col) {
    CUDA_KERNEL_LOOP(index, n) {
        int w_out = index % width_col;

        luaL_argcheck(L, batchSize == 1 || batchSize % 4 == 0, 1, "batch size should be a multiple of 4 or equal to 1");
        luaL_argcheck(L, nOutputPlane % 8 == 0, 1, "nOutputPlane should be a multiple of 8");

        // Step batch (inner loop)
        // This variable defines how many samples are processed in //, in the inner loop
        int stepBatchSize = 1;
        if (batchSize % 4 == 0) {
            stepBatchSize = 4;
        }

        // Resize output
        THCudaTensor_resize4d(output, batchSize, outputHeight, outputWidth, nOutputPlane);

        // Resize temporary columns
        THCudaTensor_resize2d(columns, kH*kW*nInputPlane, stepBatchSize*outputHeight*outputWidth);

        // Add bias first
        // TODO: replace this by more efficient, custom kernel
        long k;
        THCudaTensor *outputPlane = THCudaTensor_new();
        for(k=0; k<nOutputPlane; k++) {
