     S.emplace_back(tanh(dot(Y[t], Wxh_d) + dot(S[t-1], Whh_d) + bh_d));
   }
 
-  Expr Why = named(g.param(shape={hidden_size, output_size}, init=uniform()), "Why");
-  Expr by = named(g.param(shape={1, output_size}, init=uniform()), "by");
+  // Output linear layer before softmax.
+  Expr Why = named(g.param(shape={hidden_size, output_size},
+                           init=uniform()), "Why");
+  Expr by = named(g.param(shape={1, output_size},
+                          init=uniform()), "by");
 
   std::cerr << "Building output layer..." << std::endl;
+
+  // Softmax layer and cost function.
   std::vector<Expr> Yp;
-
   Yp.emplace_back(named(softmax_fast(dot(h0_d, Why) + by), "pred"));
   Expr cross_entropy = sum(Y[0] * log(Yp[0]), axis=1);
   for (int t = 1; t <= num_outputs; ++t) {

  std::cerr << "Target vocabulary size: " << targetVocab.Size() << std::endl;
#endif

  ExpressionGraph g = build_graph(0);

  // For the stop symbol.
  for (int t = 0; t <= num_inputs; ++t) {
    Tensor Xt({batch_size, input_size});

    float max = 1.;
    std::vector<float> values(batch_size * input_size);
    std::vector<float> classes(batch_size * output_size, 0.0);

  ExpressionGraph g(cuda_device);
  std::vector<Expr> X, Y, H, S;

  // We're including the stop symbol here.
  for (int t = 0; t <= num_inputs; ++t) {
    std::stringstream ss;
    ss << "X" << t;
    X.emplace_back(named(g.input(shape={batch_size, input_size}), ss.str()));
  }

  // We're including the stop symbol here.
  for (int t = 0; t <= num_outputs; ++t) {
    std::stringstream ss;
    ss << "Y" << t;
    Y.emplace_back(named(g.input(shape={batch_size, output_size}), ss.str()));
  }

  // Source embeddings.
  Expr E = named(g.param(shape={input_size, embedding_size},
                         init=uniform()), "E");

  // Source RNN parameters.
  Expr Wxh = named(g.param(shape={embedding_size, hidden_size},
                   init=uniform()), "Wxh");
  Expr Whh = named(g.param(shape={hidden_size, hidden_size},
                   init=uniform()), "Whh");
  Expr bh = named(g.param(shape={1, hidden_size},
                  init=uniform()), "bh");
  Expr h0 = named(g.param(shape={1, hidden_size},
                  init=uniform()), "h0");

  std::cerr << "Building encoder RNN..." << std::endl;
  H.emplace_back(tanh(dot(dot(X[0], E), Wxh) + dot(h0, Whh) + bh));
  for (int t = 1; t <= num_inputs; ++t) {
    H.emplace_back(tanh(dot(dot(X[t], E), Wxh) + dot(H[t-1], Whh) + bh));
  }

  // Target RNN parameters.
  Expr Wxh_d = named(g.param(shape={output_size, hidden_size},
                     init=uniform()), "Wxh_d");
  Expr Whh_d = named(g.param(shape={hidden_size, hidden_size},
                     init=uniform()), "Whh_d");
  Expr bh_d = named(g.param(shape={1, hidden_size},
                    init=uniform()), "bh_d");

  std::cerr << "Building decoder RNN..." << std::endl;
  auto h0_d = H[num_inputs];
