polish gpt2 example

asyml · Apr 20, 2019 · 3efd3e0 · 3efd3e0
1 parent 25da80d
commit 3efd3e0
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/examples/gpt-2/gpt2_generate_main.py b/examples/gpt-2/gpt2_generate_main.py
@@ -112,10 +112,10 @@ def main(_):
 
     pos_embedder = tx.modules.PositionEmbedder(
         position_size=gpt2_config.position_size,
-        hparams=gpt2_config.pos_embed
-    )
+        hparams=gpt2_config.pos_embed)
 
     def _embedding_fn(x, y):
+        # `x` is token ids, `y` is time steps
         return word_embedder(x) + pos_embedder(y)
 
     helper = tx.modules.TopKSampleEmbeddingHelper(