 #undef _GLIBCXX_USE_INT128
 
 #include <iostream>
+#include <map>
 #include <vector>
+#include <stdio.h>
 #include <cublas_v2.h>
 #include <thrust/device_vector.h>
 #include <thrust/host_vector.h>

    string host = buf;
    host = host.substr(0, host.size() - 1);
    strcpy(buf, host.c_str());

    int devCount;
    int deviceNum=-1;
    CUDA_CHECK(cudaGetDeviceCount(&devCount));

    if (commRank == 0)
    {
        map<string, vector<int> > hosts;
        map<string, int> devCounts;
        MPI_Status stat;
        MPI_Request req;

        hosts[buf].push_back(0);
        devCounts[buf] = devCount;
        for (int i = 1; i < commSize; ++i)
        {
            MPI_Recv(buf, 1024, MPI_CHAR, i, 0, MPI_COMM_WORLD, &stat);

        MPI_Send(&devCount, 1, MPI_INT, 0, 0, MPI_COMM_WORLD);
        MPI_Recv(&deviceNum, 1, MPI_INT, 0, 0, MPI_COMM_WORLD, &stat);
    }
    MPI_Barrier(MPI_COMM_WORLD);
#endif    
    CUDA_CHECK(cudaSetDevice(deviceNum));
}

/** One epoch on the GPU, dense variant
 */
