add rotary emb option to train dalle (lucidrains#348)

Author: rom1504

train_dalle.py

@@ -132,6 +132,8 @@
 
 model_group.add_argument('--shift_tokens', default = False, type = bool, help = 'Use the shift tokens feature')
 
+model_group.add_argument('--rotary_emb', default = False, type = bool, help = 'Use rotary embeddings')
+
 args = parser.parse_args()
 
 # helpers
@@ -186,6 +188,7 @@ def cp_path_to_dir(cp_path, tag):
 ATTN_DROPOUT = args.attn_dropout
 STABLE = args.stable_softmax
 SHIFT_TOKENS = args.shift_tokens
+ROTARY_EMB = args.rotary_emb
 
 ATTN_TYPES = tuple(args.attn_types.split(','))
 
@@ -299,6 +302,7 @@ def cp_path_to_dir(cp_path, tag):
         attn_dropout=ATTN_DROPOUT,
         stable=STABLE,
         shift_tokens=SHIFT_TOKENS,
+        rotary_emb=ROTARY_EMB,
     )
     resume_epoch = 0