{
  "_class_name": "LightningDiT2DModel",
  "depth": 28,
  "hidden_size": 1152,
  "in_channels": 32,
  "input_size": 16,
  "num_classes": 1000,
  "num_heads": 16,
  "patch_size": 1,
  "use_qknorm": true,
  "use_rmsnorm": true,
  "use_rope": true,
  "use_swiglu": true,
  "wo_shift": false
}