Update README.md

README.md

@@ -52,4 +52,6 @@ Paper: [Transfusion: Predict the Next Token and Diffuse Images with One Multi-Mo
 
 Dataset: [ImageNet](https://image-net.org/), [COCO](https://cocodataset.org/), [FFHQ](https://github.com/NVlabs/ffhq-dataset)
 
-Base Code: [lavinal712/AutoencoderKL](https://github.com/lavinal712/AutoencoderKL)
+Base Code: [lavinal712/AutoencoderKL](https://github.com/lavinal712/AutoencoderKL)
+
+Training Code: [lavinal712/AutoencoderKL/tree/transfusion_vae](https://github.com/lavinal712/AutoencoderKL/tree/transfusion_vae)