diff --git a/Diffusion_UKAN/Diffusion/Train.py b/Diffusion_UKAN/Diffusion/Train.py
index d112135..57b7460 100644
--- a/Diffusion_UKAN/Diffusion/Train.py
+++ b/Diffusion_UKAN/Diffusion/Train.py
@@ -9,7 +9,7 @@
 # from torchvision.datasets import CIFAR10
 from torchvision.utils import save_image
 from Diffusion import GaussianDiffusionSampler, GaussianDiffusionTrainer
-from Diffusion.UNet import UNet, UNe    t_Baseline
+from Diffusion.UNet import UNet, UNet_Baseline
 from Diffusion.Model_ConvKan import UNet_ConvKan
 from Diffusion.Model_UMLP import UMLP
 from Diffusion.Model_UKAN_Hybrid import UKan_Hybrid