 
     // Allocate memory on the device
 
-    T* tmp_loss;
-    T* tmp_error;
-    cuda_check(cudaMalloc((void**)&tmp_loss, numBlocks * sizeof(T)));
-    cuda_check(cudaMalloc((void**)&tmp_error, numBlocks * sizeof(T)));
+    T* tmp_gpu;
+    cuda_check(cudaMalloc((void**)&tmp_gpu, 2 * numBlocks * sizeof(T)));
+
+    T* tmp_loss = tmp_gpu;
+    T* tmp_error = tmp_gpu + numBlocks;
 
     // Run the first reduction on GPU
 

    }

    if(s > 1){
        T* host_data = new T[s];

        cuda_check(cudaMemcpy(host_data, tmp_loss, s * sizeof(T), cudaMemcpyDeviceToHost));

        for (size_t i = 0; i < s; i++) {
            loss += host_data[i];
        }

        cuda_check(cudaMemcpy(host_data, tmp_error, s * sizeof(T), cudaMemcpyDeviceToHost));

        for (size_t i = 0; i < s; i++) {
            error += host_data[i];
        }

        delete[] host_data;

        cuda_check(cudaMemcpy(&error, tmp_error, 1 * sizeof(T), cudaMemcpyDeviceToHost));
    }

    cuda_check(cudaFree(tmp_gpu));

    return std::make_pair(loss, error);
}
