 namespace marian {
 
   void TrainingLoop(Config& options,
-                    Ptr<data::BatchGenerator<data::Corpus>> batchGenerator,
-                    std::function<float(Ptr<data::CorpusBatch>)> update,
-                    std::function<void(const std::string&)> save) {
+                    Ptr<data::BatchGenerator<data::Corpus>> batchGenerator) {
+ 
+ 
+    auto dimVocabs = options.get<std::vector<int>>("dim-vocabs");
+    int dimEmb = options.get<int>("dim-emb");
+    int dimRnn = options.get<int>("dim-rnn");
+    int dimBatch = options.get<int>("mini-batch");
+    int dimMaxiBatch = options.get<int>("maxi-batch");
+
+    std::vector<int> dims = {
+      dimVocabs[0], dimEmb, dimRnn,
+      dimVocabs[1], dimEmb, dimRnn,
+      dimBatch
+    };
+
+    auto devices = options.get<std::vector<int>>("device");
+    
+    std::mutex dispMutex;
+    
+    auto save = [](const std::string& name) {
+      LOG(info) << "Saving parameters to " << name;
+      //nematus->save(graph, name);
+    };
+
+    ThreadPool pool(devices.size(), devices.size());
+ 
     boost::timer::cpu_timer timer;
 
     size_t epochs = 1;

  int dimRnn = options.get<int>("dim-rnn");
  int dimBatch = options.get<int>("mini-batch");
  int dimMaxiBatch = options.get<int>("maxi-batch");

  std::vector<int> dims = {
    dimVocabs[0], dimEmb, dimRnn,
    dimVocabs[1], dimEmb, dimRnn,
    dimBatch
  };

  auto devices = options.get<std::vector<int>>("device");
  int device = devices[0];

  auto graph = New<ExpressionGraph>();
  graph->setDevice(device);

  auto nematus = New<Nematus>(dims);

  if(options.has("init")) {
    LOG(info) << "Loading parameters from " << options.get<std::string>("init");
    nematus->load(graph, options.get<std::string>("init"));
  }

  graph->reserveWorkspaceMB(options.get<size_t>("workspace"));

  Ptr<ClipperBase> clipper = nullptr;

  float clipNorm = options.get<double>("clip-norm");
  float lrate = options.get<double>("lrate");
  if(clipNorm > 0)
    clipper = Clipper<Norm>(clipNorm);
  auto opt = Optimizer<Adam>(lrate, clip=clipper);

  auto update = [graph, opt, nematus](Ptr<CorpusBatch> batch) -> float {
    auto costNode = nematus->construct(graph, batch);
    opt->update(graph);
    return costNode->scalar();
  };

  auto save = [graph, nematus](const std::string& name) {
    LOG(info) << "Saving parameters to " << name;
    nematus->save(graph, name);
  };

  auto trainSets = options.get<std::vector<std::string>>("trainsets");
  auto vocabs = options.get<std::vector<std::string>>("vocabs");
  size_t maxSentenceLength = options.get<size_t>("max-length");

  auto corpus = New<Corpus>(trainSets, vocabs, dimVocabs, maxSentenceLength);
  auto bg = New<BatchGenerator<Corpus>>(corpus, dimBatch, dimMaxiBatch);

  TrainingLoop(options, bg, update, save);

  return 0;
}

      size_t wordsDisp = 0;

      while(*batchGenerator) {
        
        for(auto device : devices) {
        
          auto update = [&, device](Ptr<data::CorpusBatch> batch) {
            thread_local Ptr<ExpressionGraph> graph;
            thread_local Ptr<Nematus> nematus;
            thread_local Ptr<OptimizerBase> opt;
            
            if(!graph) {
              std::lock_guard<std::mutex> guard(dispMutex);
              graph = New<ExpressionGraph>();
              graph->setDevice(device);
              graph->reserveWorkspaceMB(options.get<size_t>("workspace"));
            }
            
            if(!nematus) {
              nematus = New<Nematus>(dims);
              if(options.has("init")) {
                LOG(info) << "Loading parameters from " << options.get<std::string>("init");
                nematus->load(graph, options.get<std::string>("init"));
              }
            }
          
            if(!opt) {
              Ptr<ClipperBase> clipper = nullptr;
              float clipNorm = options.get<double>("clip-norm");
              float lrate = options.get<double>("lrate");
              if(clipNorm > 0)
                clipper = Clipper<Norm>(clipNorm);
              opt = Optimizer<Adam>(lrate, keywords::clip=clipper);
            }
            
            auto costNode = nematus->construct(graph, batch);
            opt->update(graph);
            float cost =  costNode->scalar();
            
            std::lock_guard<std::mutex> guard(dispMutex);
        
            costSum += cost;
            samples += batch->size();
            wordsDisp += batch->words();
            batches++;
            //if(options.get<size_t>("after-batches")
            //   && batches >= options.get<size_t>("after-batches"))
            //  break;
    
            if(batches % options.get<size_t>("disp-freq") == 0) {
              std::stringstream ss;
              ss << "Ep. " << epochs
                 << " : Up. " << batches
                 << " : Sen. " << samples
                 << " : Cost " << std::fixed << std::setprecision(2)
                               << costSum / options.get<size_t>("disp-freq")
                 << " : Time " << timer.format(2, "%ws");
    
              float seconds = std::stof(timer.format(5, "%w"));
              float wps = wordsDisp /   (float)seconds;
    
              ss << " : " << std::fixed << std::setprecision(2)
                 << wps << " words/s";
    
              LOG(info) << ss.str();
    
              timer.start();
              costSum = 0;
              wordsDisp = 0;
            }
    
            if(batches % options.get<size_t>("save-freq") == 0) {
              if(options.get<bool>("overwrite"))
                save(options.get<std::string>("model") + ".npz");
              else
                save(options.get<std::string>("model") + "." + std::to_string(batches) + ".npz");
            }
            
          };
          
          auto batch = batchGenerator->next();
          pool.enqueue(update, batch);
        }
      }
      epochs++;
