Fix DiT sampling

Apply the PixArt fix here as well.

Move models to checkpoints folder as that makes more sense in this case.

DiT/model.py

@@ -158,6 +158,7 @@ class DiT(nn.Module):
 		class_dropout_prob=0.1,
 		num_classes=1000,
 		learn_sigma=True,
+		**kwargs,
 	):
 		super().__init__()
 		self.learn_sigma = learn_sigma