   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);
+ 
+  TrainingLoop(options, bg);
 
   return 0;
 }

      size_t wordsDisp = 0;

      while(*batchGenerator) {
        auto batch = batchGenerator->next();

        float cost = update(batch);

        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");
        }
      }
      epochs++;

#include <boost/timer/timer.hpp>
#include <boost/chrono.hpp>
#include <boost/program_options.hpp>
#include <thread>
#include <chrono>
#include <mutex>

#include "marian.h"
#include "optimizers/optimizers.h"
