         return;
     }
 
+    size_t gridSize  = 64;
+    size_t blockSize = 64;
+
+    if (n <= 100) {
+        gridSize  = 1;
+        blockSize = 64;
+    } else if(n <= 1000){
+        gridSize  = 8;
+        blockSize = 64;
+    } else if(n <= 10000){
+        gridSize  = 16;
+        blockSize = 64;
+    } else if(n <= 100000){
+        gridSize  = 32;
+        blockSize = 64;
+    }
+
     // Allocate room for the states
     curandState* states;
-    cuda_check(cudaMalloc((void**)&states, 64 * 64 * sizeof(curandState)));
+    cuda_check(cudaMalloc((void**)&states, gridSize * blockSize * sizeof(curandState)));
 
     // Initialize the seeds
-    setup_kernel<<<64, 64>>>(states, seed);
+    setup_kernel<<<gridSize, blockSize>>>(states, seed);
     cudaDeviceSynchronize();
 
     // Compute the dropout mask
     if (alpha == 1.0f) {
-        inv_dropout_kernel1<double><<<64,64>>>(states, n, p, x, incx);
+        inv_dropout_kernel1<double><<<gridSize, blockSize>>>(states, n, p, x, incx);
     } else {
-        inv_dropout_kernel<double><<<64,64>>>(states, n, p, alpha, x, incx);
+        inv_dropout_kernel<double><<<gridSize, blockSize>>>(states, n, p, alpha, x, incx);
     }
 
     // Free the states

        return;
    }

    // Allocate room for the states
    curandState* states;
    cuda_check(cudaMalloc((void**)&states, 64 * 64 * sizeof(curandState)));

    // Initialize the seeds
    setup_kernel<<<64, 64>>>(states, seed);
    cudaDeviceSynchronize();

    // Compute the dropout mask
    if (alpha == 1.0f) {
        dropout_kernel1<float><<<64,64>>>(states, n, p, x, incx);
    } else {
        dropout_kernel<float><<<64,64>>>(states, n, p, alpha, x, incx);
    }

    // Free the states

        return;
    }

    size_t gridSize  = 64;
    size_t blockSize = 64;

    if (n <= 100) {
        gridSize  = 1;
        blockSize = 64;
    } else if(n <= 1000){
        gridSize  = 8;
        blockSize = 64;
    } else if(n <= 10000){
        gridSize  = 16;
        blockSize = 64;
    } else if(n <= 100000){
        gridSize  = 32;
        blockSize = 64;
    }

    // Allocate room for the states
    curandState* states;
    cuda_check(cudaMalloc((void**)&states, gridSize * blockSize * sizeof(curandState)));

    // Initialize the seeds
    setup_kernel<<<gridSize, blockSize>>>(states, seed);
    cudaDeviceSynchronize();

    // Compute the dropout mask
    if (alpha == 1.0f) {
        inv_dropout_kernel1<float><<<gridSize,blockSize>>>(states, n, p, x, incx);
    } else {
        inv_dropout_kernel<float><<<gridSize,blockSize>>>(states, n, p, alpha, x, incx);
    }

    // Free the states
