use correct decoder args

Author: lukas-blecher

pix2tex/models/transformer.py

@@ -61,6 +61,6 @@ def get_decoder(args):
                 dim=args.dim,
                 depth=args.num_layers,
                 heads=args.heads,
-                cross_attend=True
+                **args.decoder_args
             )),
         pad_value=args.pad_token)

pix2tex/train.py

@@ -70,7 +70,7 @@ def save_models(e, step=0):
                         max_bleu, max_token_acc = bleu_score, token_accuracy
                         save_models(e, step=i)
             if (e+1) % args.save_freq == 0:
-                save_models(e)
+                save_models(e, step=len(dataloader))
             if args.wandb:
                 wandb.log({'train/epoch': e+1})
     except KeyboardInterrupt: