Update.

diff --git a/main.py b/main.py
index df46652..74e70b2 100755 (executable)
--- a/main.py
+++ b/main.py
@@ -346,6 +346,7 @@ default_model_args = {
         "nb_blocks": 6,
     },
     "37M": {
+        "attention": "mha",
         "dim_model": 512,
         "dim_keys": 64,
         "dim_hidden": 2048,