 
 void imt2col(const float* data_im, const int channels,
         const int height, const int width, const int ksize, const int pad,
-        const int stride, float* data_col) {
+        const int stride, const int batch, float* data_col) {
     // We are going to launch channels * height_col * width_col kernels, each
     // kernel responsible for copying a single-channel grid.
     int height_col = (height + 2 * pad - ksize) / stride + 1;
     int width_col = (width + 2 * pad - ksize) / stride + 1;
     int num_kernels = channels * height_col * width_col;
     // Launch
-    imt2col_kernel <<<GET_BLOCKS(num_kernels), CUDA_NUM_THREADS>>> (
-        num_kernels, data_im, height, width, ksize, 
-        pad, stride, channels,
-        height_col, width_col, data_col
-    );
+    for (int bidx = 0; bidx < batch; bidx++) {
+        imt2col_kernel <<<GET_BLOCKS(num_kernels), CUDA_NUM_THREADS>>> (
+            num_kernels, data_im, height, width, ksize, 
+            pad, stride, channels,
+            height_col, width_col, bidx, batch,
+            data_col
+        );
+    }
 }
 
 static int cunn_SpatialConvolutionMM_BHWD_updateOutput(lua_State *L) {

        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++) {
