maybe you def prepare_data() does not working ? like the code " x_t, noise = diffusion.forward_diffusion(target_image, t)" diffusion function not normal operation, then x_t variable is not a good values. Lead to later matrix size mapping error.