diff --git a/training_tpu.py b/training_tpu.py
index 705d467..a7cbb11 100644
--- a/training_tpu.py
+++ b/training_tpu.py
@@ -983,6 +983,7 @@ def main(argv):
             "temb": (),
             "textcontext": (77, 768)
         },
+        "arguments": args
     }
     
     text_encoders = defaultTextEncodeModel()