Update modeling_ndlinear_dit.py

modeling_ndlinear_dit.py

@@ -8,11 +8,11 @@ from mlp import NdMlp
 from ndlinear import NdLinear
 from models_hf import DiT, DiTConfig
 
-class NdLinearDiTConfig(DiTConfig):
+class DiTConfig(PretrainedConfig):
     model_type = "ndlinear_dit"
 
-class NdLinearDiT(DiT):
-    config_class = NdLinearDiTConfig
+class DiT(PreTrainedModel):
+    config_class = DiTConfig
 
-__all__ = ["NdLinearDiT", "NdLinearDiTConfig"]
+__all__ = ["DiT", "DiTConfig"]