Upload folder using huggingface_hub

703f61c verified 11 months ago

233 kB

	Using devices [TpuDevice(id=0, process_index=0, coords=(0,0,0), core_on_chip=0), TpuDevice(id=1, process_index=0, coords=(1,0,0), core_on_chip=0), TpuDevice(id=2, process_index=0, coords=(0,1,0), core_on_chip=0), TpuDevice(id=3, process_index=0, coords=(1,1,0), core_on_chip=0)]
	Device count 4
	Global device count 4
	Global Batch: 512
	Node Batch: 512
	Device Batch: 128
	Loading dataset
	Loading dataset
	creating model
	beta1: 0.9
	beta2: 0.999
	bootstrap_cfg: 1
	bootstrap_dt_bias: 0
	bootstrap_ema: 1
	bootstrap_every: 8
	cfg_scale: 1.5
	class_dropout_prob: 0.1
	denoise_timesteps: 128
	depth: 12
	dropout: 0.0
	dt_sampling: uniform
	hidden_size: 768
	lr: 0.0001
	mlp_ratio: 4
	num_classes: 1000
	num_heads: 12
	patch_size: 2
	sharding: dp
	t_sampling: discrete-dt
	target_update_rate: 0.999
	train_type: shortcut
	use_cosine: 0
	use_ema: 0
	use_stable_vae: 1
	warmup: 0
	weight_decay: 0.1

	Total devices TPU_0(process=0,(0,0,0,0))
	Initializing encoder.
	Incoming encoder shape (1, 256, 256, 3)
	Encoder layer (1, 256, 256, 128)
	doing downsample
	Encoder layer (1, 128, 128, 128)
	doing downsample
	Encoder layer (1, 64, 64, 256)
	doing downsample
	Encoder layer (1, 32, 32, 512)
	Encoder layer (1, 32, 32, 512)
	Encoder layer final (1, 32, 32, 512)
	Encoder layer final (1, 32, 32, 512)
	Final embeddings are size (1, 32, 32, 8)
	After quant (1, 32, 32, 4)
	encode finished
	Decoder incoming shape (1, 32, 32, 4)
	Decoder input (1, 32, 32, 512)
	Mid Block Decoder layer (1, 32, 32, 512)
	Mid Block Decoder layer (1, 32, 32, 512)
	Decoder layer (1, 64, 64, 512)
	Decoder layer (1, 128, 128, 512)
	Decoder layer (1, 256, 256, 256)
	Decoder layer (1, 256, 256, 128)
	Total num of VQVAE parameters: 67565323
	Disc shape (1, 128, 128, 128)
	Disc shape (1, 64, 64, 256)
	Disc shape (1, 32, 32, 512)
	Disc shape (1, 16, 16, 512)
	Disc shape (1, 8, 8, 512)
	Disc shape (1, 4, 4, 512)
	Total num of Discriminator parameters: 23998017
	Loaded checkpoint from 11024065 seconds ago.
	Loaded model with step 500001
	┌──────────────────────────────────────────────────────────────────────────────┐
	│ TPU 0 │
	├──────────────────────────────────────────────────────────────────────────────┤
	│ TPU 1 │
	├──────────────────────────────────────────────────────────────────────────────┤
	│ TPU 2 │
	├──────────────────────────────────────────────────────────────────────────────┤
	│ TPU 3 │
	└──────────────────────────────────────────────────────────────────────────────┘
	returning model
	model done
	Input to vae (4, 1, 256, 256, 3)
	encode image shape (1, 256, 256, 3)
	Initializing encoder.
	Incoming encoder shape (1, 256, 256, 3)
	Encoder layer (1, 256, 256, 128)
	doing downsample
	Encoder layer (1, 128, 128, 128)
	doing downsample
	Encoder layer (1, 64, 64, 256)
	doing downsample
	Encoder layer (1, 32, 32, 512)
	Encoder layer (1, 32, 32, 512)
	Encoder layer final (1, 32, 32, 512)
	Encoder layer final (1, 32, 32, 512)
	Final embeddings are size (1, 32, 32, 8)
	After quant (1, 32, 32, 4)
	output example shape (4, 1, 32, 32, 4)
	Test data shape (4, 256, 256, 3)
	x shape (4, 1, 256, 256, 3)
	encoded shape (4, 1, 32, 32, 4)
	z_vectors shape (1, 32, 32, 4)
	Decoder incoming shape (1, 32, 32, 4)
	Decoder input (1, 32, 32, 512)
	Mid Block Decoder layer (1, 32, 32, 512)
	Mid Block Decoder layer (1, 32, 32, 512)
	Decoder layer (1, 64, 64, 512)
	Decoder layer (1, 128, 128, 512)
	Decoder layer (1, 256, 256, 256)
	Decoder layer (1, 256, 256, 128)
	image shape (4, 1, 256, 256, 3)
	decoded img shape (256, 256, 3)
	obs shape (4, 32, 32, 4)
	DiT: Input of shape (4, 32, 32, 4) dtype float32
	DiT: After patch embed, shape is (4, 256, 768) dtype bfloat16
	DiT: Patch Embed of shape (4, 256, 768) dtype bfloat16
	DiT: Conditioning of shape (1, 768) dtype float32

	[3m DiT Summary [0m
	┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓
	┃[1m [0m[1mpath [0m[1m [0m┃[1m [0m[1mmodule [0m[1m [0m┃[1m [0m[1minputs [0m[1m [0m┃[1m [0m[1moutputs [0m[1m [0m┃[1m [0m[1mparams [0m[1m [0m┃
	┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┩
	│ │ DiT │ - [2mfloat32[0m[4,32,32,4] │ [2mbfloat16[0m[4,32,32,4] │ │
	│ │ │ - [2mfloat32[0m[1] │ │ │
	│ │ │ - [2mfloat32[0m[1] │ │ │
	│ │ │ - [2mint32[0m[1] │ │ │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ PatchEmbed_0 │ PatchEmbed │ [2mfloat32[0m[4,32,32,4] │ [2mbfloat16[0m[4,256,768] │ │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ PatchEmbed_0/Conv_0 │ Conv │ [2mfloat32[0m[4,32,32,4] │ [2mbfloat16[0m[4,16,16,768] │ bias: [2mfloat32[0m[768] │
	│ │ │ │ │ kernel: [2mfloat32[0m[2,2,4,768] │
	│ │ │ │ │ │
	│ │ │ │ │ [1m13,056 [0m[1;2m(52.2 KB)[0m │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ TimestepEmbedder_0 │ TimestepEmbedder │ [2mfloat32[0m[1] │ [2mfloat32[0m[1,768] │ │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ TimestepEmbedder_0/Dense_0 │ Dense │ [2mbfloat16[0m[1,256] │ [2mbfloat16[0m[1,768] │ bias: [2mfloat32[0m[768] │
	│ │ │ │ │ kernel: [2mfloat32[0m[256,768] │
	│ │ │ │ │ │
	│ │ │ │ │ [1m197,376 [0m[1;2m(789.5 KB)[0m │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ TimestepEmbedder_0/Dense_1 │ Dense │ [2mbfloat16[0m[1,768] │ [2mfloat32[0m[1,768] │ bias: [2mfloat32[0m[768] │
	│ │ │ │ │ kernel: [2mfloat32[0m[768,768] │
	│ │ │ │ │ │
	│ │ │ │ │ [1m590,592 [0m[1;2m(2.4 MB)[0m │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ TimestepEmbedder_1 │ TimestepEmbedder │ [2mfloat32[0m[1] │ [2mfloat32[0m[1,768] │ │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ TimestepEmbedder_1/Dense_0 │ Dense │ [2mbfloat16[0m[1,256] │ [2mbfloat16[0m[1,768] │ bias: [2mfloat32[0m[768] │
	│ │ │ │ │ kernel: [2mfloat32[0m[256,768] │
	│ │ │ │ │ │
	│ │ │ │ │ [1m197,376 [0m[1;2m(789.5 KB)[0m │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ TimestepEmbedder_1/Dense_1 │ Dense │ [2mbfloat16[0m[1,768] │ [2mfloat32[0m[1,768] │ bias: [2mfloat32[0m[768] │
	│ │ │ │ │ kernel: [2mfloat32[0m[768,768] │
	│ │ │ │ │ │
	│ │ │ │ │ [1m590,592 [0m[1;2m(2.4 MB)[0m │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ LabelEmbedder_0 │ LabelEmbedder │ [2mint32[0m[1] │ [2mbfloat16[0m[1,768] │ │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ LabelEmbedder_0/Embed_0 │ Embed │ [2mint32[0m[1] │ [2mbfloat16[0m[1,768] │ embedding: [2mfloat32[0m[1001,768] │
	│ │ │ │ │ │
	│ │ │ │ │ [1m768,768 [0m[1;2m(3.1 MB)[0m │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_0 │ DiTBlock │ - [2mbfloat16[0m[4,256,768] │ [2mbfloat16[0m[4,256,768] │ │
	│ │ │ - [2mfloat32[0m[1,768] │ │ │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_0/Dense_0 │ Dense │ [2mfloat32[0m[1,768] │ [2mbfloat16[0m[1,4608] │ bias: [2mfloat32[0m[4608] │
	│ │ │ │ │ kernel: [2mfloat32[0m[768,4608] │
	│ │ │ │ │ │
	│ │ │ │ │ [1m3,543,552 [0m[1;2m(14.2 MB)[0m │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_0/LayerNorm_0 │ LayerNorm │ [2mbfloat16[0m[4,256,768] │ [2mbfloat16[0m[4,256,768] │ │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_0/Dense_1 │ Dense │ [2mbfloat16[0m[4,256,768] │ [2mbfloat16[0m[4,256,768] │ bias: [2mfloat32[0m[768] │
	│ │ │ │ │ kernel: [2mfloat32[0m[768,768] │
	│ │ │ │ │ │
	│ │ │ │ │ [1m590,592 [0m[1;2m(2.4 MB)[0m │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_0/Dense_2 │ Dense │ [2mbfloat16[0m[4,256,768] │ [2mbfloat16[0m[4,256,768] │ bias: [2mfloat32[0m[768] │
	│ │ │ │ │ kernel: [2mfloat32[0m[768,768] │
	│ │ │ │ │ │
	│ │ │ │ │ [1m590,592 [0m[1;2m(2.4 MB)[0m │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_0/Dense_3 │ Dense │ [2mbfloat16[0m[4,256,768] │ [2mbfloat16[0m[4,256,768] │ bias: [2mfloat32[0m[768] │
	│ │ │ │ │ kernel: [2mfloat32[0m[768,768] │
	│ │ │ │ │ │
	│ │ │ │ │ [1m590,592 [0m[1;2m(2.4 MB)[0m │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_0/Dense_4 │ Dense │ [2mfloat32[0m[4,256,768] │ [2mbfloat16[0m[4,256,768] │ bias: [2mfloat32[0m[768] │
	│ │ │ │ │ kernel: [2mfloat32[0m[768,768] │
	│ │ │ │ │ │
	│ │ │ │ │ [1m590,592 [0m[1;2m(2.4 MB)[0m │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_0/LayerNorm_1 │ LayerNorm │ [2mbfloat16[0m[4,256,768] │ [2mbfloat16[0m[4,256,768] │ │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_0/MlpBlock_0 │ MlpBlock │ [2mbfloat16[0m[4,256,768] │ [2mbfloat16[0m[4,256,768] │ │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_0/MlpBlock_0/Dense_0 │ Dense │ [2mbfloat16[0m[4,256,768] │ [2mbfloat16[0m[4,256,3072] │ bias: [2mfloat32[0m[3072] │
	│ │ │ │ │ kernel: [2mfloat32[0m[768,3072] │
	│ │ │ │ │ │
	│ │ │ │ │ [1m2,362,368 [0m[1;2m(9.4 MB)[0m │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_0/MlpBlock_0/Dropout_0 │ Dropout │ [2mbfloat16[0m[4,256,3072] │ [2mbfloat16[0m[4,256,3072] │ │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_0/MlpBlock_0/Dense_1 │ Dense │ [2mbfloat16[0m[4,256,3072] │ [2mbfloat16[0m[4,256,768] │ bias: [2mfloat32[0m[768] │
	│ │ │ │ │ kernel: [2mfloat32[0m[3072,768] │
	│ │ │ │ │ │
	│ │ │ │ │ [1m2,360,064 [0m[1;2m(9.4 MB)[0m │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_0/MlpBlock_0/Dropout_1 │ Dropout │ [2mbfloat16[0m[4,256,768] │ [2mbfloat16[0m[4,256,768] │ │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_1 │ DiTBlock │ - [2mbfloat16[0m[4,256,768] │ [2mbfloat16[0m[4,256,768] │ │
	│ │ │ - [2mfloat32[0m[1,768] │ │ │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_1/Dense_0 │ Dense │ [2mfloat32[0m[1,768] │ [2mbfloat16[0m[1,4608] │ bias: [2mfloat32[0m[4608] │
	│ │ │ │ │ kernel: [2mfloat32[0m[768,4608] │
	│ │ │ │ │ │
	│ │ │ │ │ [1m3,543,552 [0m[1;2m(14.2 MB)[0m │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_1/LayerNorm_0 │ LayerNorm │ [2mbfloat16[0m[4,256,768] │ [2mbfloat16[0m[4,256,768] │ │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_1/Dense_1 │ Dense │ [2mbfloat16[0m[4,256,768] │ [2mbfloat16[0m[4,256,768] │ bias: [2mfloat32[0m[768] │
	│ │ │ │ │ kernel: [2mfloat32[0m[768,768] │
	│ │ │ │ │ │
	│ │ │ │ │ [1m590,592 [0m[1;2m(2.4 MB)[0m │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_1/Dense_2 │ Dense │ [2mbfloat16[0m[4,256,768] │ [2mbfloat16[0m[4,256,768] │ bias: [2mfloat32[0m[768] │
	│ │ │ │ │ kernel: [2mfloat32[0m[768,768] │
	│ │ │ │ │ │
	│ │ │ │ │ [1m590,592 [0m[1;2m(2.4 MB)[0m │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_1/Dense_3 │ Dense │ [2mbfloat16[0m[4,256,768] │ [2mbfloat16[0m[4,256,768] │ bias: [2mfloat32[0m[768] │
	│ │ │ │ │ kernel: [2mfloat32[0m[768,768] │
	│ │ │ │ │ │
	│ │ │ │ │ [1m590,592 [0m[1;2m(2.4 MB)[0m │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_1/Dense_4 │ Dense │ [2mfloat32[0m[4,256,768] │ [2mbfloat16[0m[4,256,768] │ bias: [2mfloat32[0m[768] │
	│ │ │ │ │ kernel: [2mfloat32[0m[768,768] │
	│ │ │ │ │ │
	│ │ │ │ │ [1m590,592 [0m[1;2m(2.4 MB)[0m │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_1/LayerNorm_1 │ LayerNorm │ [2mbfloat16[0m[4,256,768] │ [2mbfloat16[0m[4,256,768] │ │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_1/MlpBlock_0 │ MlpBlock │ [2mbfloat16[0m[4,256,768] │ [2mbfloat16[0m[4,256,768] │ │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_1/MlpBlock_0/Dense_0 │ Dense │ [2mbfloat16[0m[4,256,768] │ [2mbfloat16[0m[4,256,3072] │ bias: [2mfloat32[0m[3072] │
	│ │ │ │ │ kernel: [2mfloat32[0m[768,3072] │
	│ │ │ │ │ │
	│ │ │ │ │ [1m2,362,368 [0m[1;2m(9.4 MB)[0m │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_1/MlpBlock_0/Dropout_0 │ Dropout │ [2mbfloat16[0m[4,256,3072] │ [2mbfloat16[0m[4,256,3072] │ │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_1/MlpBlock_0/Dense_1 │ Dense │ [2mbfloat16[0m[4,256,3072] │ [2mbfloat16[0m[4,256,768] │ bias: [2mfloat32[0m[768] │
	│ │ │ │ │ kernel: [2mfloat32[0m[3072,768] │
	│ │ │ │ │ │
	│ │ │ │ │ [1m2,360,064 [0m[1;2m(9.4 MB)[0m │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_1/MlpBlock_0/Dropout_1 │ Dropout │ [2mbfloat16[0m[4,256,768] │ [2mbfloat16[0m[4,256,768] │ │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_2 │ DiTBlock │ - [2mbfloat16[0m[4,256,768] │ [2mbfloat16[0m[4,256,768] │ │
	│ │ │ - [2mfloat32[0m[1,768] │ │ │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_2/Dense_0 │ Dense │ [2mfloat32[0m[1,768] │ [2mbfloat16[0m[1,4608] │ bias: [2mfloat32[0m[4608] │
	│ │ │ │ │ kernel: [2mfloat32[0m[768,4608] │
	│ │ │ │ │ │
	│ │ │ │ │ [1m3,543,552 [0m[1;2m(14.2 MB)[0m │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_2/LayerNorm_0 │ LayerNorm │ [2mbfloat16[0m[4,256,768] │ [2mbfloat16[0m[4,256,768] │ │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_2/Dense_1 │ Dense │ [2mbfloat16[0m[4,256,768] │ [2mbfloat16[0m[4,256,768] │ bias: [2mfloat32[0m[768] │
	│ │ │ │ │ kernel: [2mfloat32[0m[768,768] │
	│ │ │ │ │ │
	│ │ │ │ │ [1m590,592 [0m[1;2m(2.4 MB)[0m │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_2/Dense_2 │ Dense │ [2mbfloat16[0m[4,256,768] │ [2mbfloat16[0m[4,256,768] │ bias: [2mfloat32[0m[768] │
	│ │ │ │ │ kernel: [2mfloat32[0m[768,768] │
	│ │ │ │ │ │
	│ │ │ │ │ [1m590,592 [0m[1;2m(2.4 MB)[0m │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_2/Dense_3 │ Dense │ [2mbfloat16[0m[4,256,768] │ [2mbfloat16[0m[4,256,768] │ bias: [2mfloat32[0m[768] │
	│ │ │ │ │ kernel: [2mfloat32[0m[768,768] │
	│ │ │ │ │ │
	│ │ │ │ │ [1m590,592 [0m[1;2m(2.4 MB)[0m │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_2/Dense_4 │ Dense │ [2mfloat32[0m[4,256,768] │ [2mbfloat16[0m[4,256,768] │ bias: [2mfloat32[0m[768] │
	│ │ │ │ │ kernel: [2mfloat32[0m[768,768] │
	│ │ │ │ │ │
	│ │ │ │ │ [1m590,592 [0m[1;2m(2.4 MB)[0m │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_2/LayerNorm_1 │ LayerNorm │ [2mbfloat16[0m[4,256,768] │ [2mbfloat16[0m[4,256,768] │ │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_2/MlpBlock_0 │ MlpBlock │ [2mbfloat16[0m[4,256,768] │ [2mbfloat16[0m[4,256,768] │ │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_2/MlpBlock_0/Dense_0 │ Dense │ [2mbfloat16[0m[4,256,768] │ [2mbfloat16[0m[4,256,3072] │ bias: [2mfloat32[0m[3072] │
	│ │ │ │ │ kernel: [2mfloat32[0m[768,3072] │
	│ │ │ │ │ │
	│ │ │ │ │ [1m2,362,368 [0m[1;2m(9.4 MB)[0m │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_2/MlpBlock_0/Dropout_0 │ Dropout │ [2mbfloat16[0m[4,256,3072] │ [2mbfloat16[0m[4,256,3072] │ │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_2/MlpBlock_0/Dense_1 │ Dense │ [2mbfloat16[0m[4,256,3072] │ [2mbfloat16[0m[4,256,768] │ bias: [2mfloat32[0m[768] │
	│ │ │ │ │ kernel: [2mfloat32[0m[3072,768] │
	│ │ │ │ │ │
	│ │ │ │ │ [1m2,360,064 [0m[1;2m(9.4 MB)[0m │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_2/MlpBlock_0/Dropout_1 │ Dropout │ [2mbfloat16[0m[4,256,768] │ [2mbfloat16[0m[4,256,768] │ │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_3 │ DiTBlock │ - [2mbfloat16[0m[4,256,768] │ [2mbfloat16[0m[4,256,768] │ │
	│ │ │ - [2mfloat32[0m[1,768] │ │ │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_3/Dense_0 │ Dense │ [2mfloat32[0m[1,768] │ [2mbfloat16[0m[1,4608] │ bias: [2mfloat32[0m[4608] │
	│ │ │ │ │ kernel: [2mfloat32[0m[768,4608] │
	│ │ │ │ │ │
	│ │ │ │ │ [1m3,543,552 [0m[1;2m(14.2 MB)[0m │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_3/LayerNorm_0 │ LayerNorm │ [2mbfloat16[0m[4,256,768] │ [2mbfloat16[0m[4,256,768] │ │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_3/Dense_1 │ Dense │ [2mbfloat16[0m[4,256,768] │ [2mbfloat16[0m[4,256,768] │ bias: [2mfloat32[0m[768] │
	│ │ │ │ │ kernel: [2mfloat32[0m[768,768] │
	│ │ │ │ │ │
	│ │ │ │ │ [1m590,592 [0m[1;2m(2.4 MB)[0m │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_3/Dense_2 │ Dense │ [2mbfloat16[0m[4,256,768] │ [2mbfloat16[0m[4,256,768] │ bias: [2mfloat32[0m[768] │
	│ │ │ │ │ kernel: [2mfloat32[0m[768,768] │
	│ │ │ │ │ │
	│ │ │ │ │ [1m590,592 [0m[1;2m(2.4 MB)[0m │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_3/Dense_3 │ Dense │ [2mbfloat16[0m[4,256,768] │ [2mbfloat16[0m[4,256,768] │ bias: [2mfloat32[0m[768] │
	│ │ │ │ │ kernel: [2mfloat32[0m[768,768] │
	│ │ │ │ │ │
	│ │ │ │ │ [1m590,592 [0m[1;2m(2.4 MB)[0m │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_3/Dense_4 │ Dense │ [2mfloat32[0m[4,256,768] │ [2mbfloat16[0m[4,256,768] │ bias: [2mfloat32[0m[768] │
	│ │ │ │ │ kernel: [2mfloat32[0m[768,768] │
	│ │ │ │ │ │
	│ │ │ │ │ [1m590,592 [0m[1;2m(2.4 MB)[0m │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_3/LayerNorm_1 │ LayerNorm │ [2mbfloat16[0m[4,256,768] │ [2mbfloat16[0m[4,256,768] │ │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_3/MlpBlock_0 │ MlpBlock │ [2mbfloat16[0m[4,256,768] │ [2mbfloat16[0m[4,256,768] │ │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_3/MlpBlock_0/Dense_0 │ Dense │ [2mbfloat16[0m[4,256,768] │ [2mbfloat16[0m[4,256,3072] │ bias: [2mfloat32[0m[3072] │
	│ │ │ │ │ kernel: [2mfloat32[0m[768,3072] │
	│ │ │ │ │ │
	│ │ │ │ │ [1m2,362,368 [0m[1;2m(9.4 MB)[0m │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_3/MlpBlock_0/Dropout_0 │ Dropout │ [2mbfloat16[0m[4,256,3072] │ [2mbfloat16[0m[4,256,3072] │ │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_3/MlpBlock_0/Dense_1 │ Dense │ [2mbfloat16[0m[4,256,3072] │ [2mbfloat16[0m[4,256,768] │ bias: [2mfloat32[0m[768] │
	│ │ │ │ │ kernel: [2mfloat32[0m[3072,768] │
	│ │ │ │ │ │
	│ │ │ │ │ [1m2,360,064 [0m[1;2m(9.4 MB)[0m │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_3/MlpBlock_0/Dropout_1 │ Dropout │ [2mbfloat16[0m[4,256,768] │ [2mbfloat16[0m[4,256,768] │ │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_4 │ DiTBlock │ - [2mbfloat16[0m[4,256,768] │ [2mbfloat16[0m[4,256,768] │ │
	│ │ │ - [2mfloat32[0m[1,768] │ │ │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_4/Dense_0 │ Dense │ [2mfloat32[0m[1,768] │ [2mbfloat16[0m[1,4608] │ bias: [2mfloat32[0m[4608] │
	│ │ │ │ │ kernel: [2mfloat32[0m[768,4608] │
	│ │ │ │ │ │
	│ │ │ │ │ [1m3,543,552 [0m[1;2m(14.2 MB)[0m │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_4/LayerNorm_0 │ LayerNorm │ [2mbfloat16[0m[4,256,768] │ [2mbfloat16[0m[4,256,768] │ │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_4/Dense_1 │ Dense │ [2mbfloat16[0m[4,256,768] │ [2mbfloat16[0m[4,256,768] │ bias: [2mfloat32[0m[768] │
	│ │ │ │ │ kernel: [2mfloat32[0m[768,768] │
	│ │ │ │ │ │
	│ │ │ │ │ [1m590,592 [0m[1;2m(2.4 MB)[0m │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_4/Dense_2 │ Dense │ [2mbfloat16[0m[4,256,768] │ [2mbfloat16[0m[4,256,768] │ bias: [2mfloat32[0m[768] │
	│ │ │ │ │ kernel: [2mfloat32[0m[768,768] │
	│ │ │ │ │ │
	│ │ │ │ │ [1m590,592 [0m[1;2m(2.4 MB)[0m │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_4/Dense_3 │ Dense │ [2mbfloat16[0m[4,256,768] │ [2mbfloat16[0m[4,256,768] │ bias: [2mfloat32[0m[768] │
	│ │ │ │ │ kernel: [2mfloat32[0m[768,768] │
	│ │ │ │ │ │
	│ │ │ │ │ [1m590,592 [0m[1;2m(2.4 MB)[0m │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_4/Dense_4 │ Dense │ [2mfloat32[0m[4,256,768] │ [2mbfloat16[0m[4,256,768] │ bias: [2mfloat32[0m[768] │
	│ │ │ │ │ kernel: [2mfloat32[0m[768,768] │
	│ │ │ │ │ │
	│ │ │ │ │ [1m590,592 [0m[1;2m(2.4 MB)[0m │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_4/LayerNorm_1 │ LayerNorm │ [2mbfloat16[0m[4,256,768] │ [2mbfloat16[0m[4,256,768] │ │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_4/MlpBlock_0 │ MlpBlock │ [2mbfloat16[0m[4,256,768] │ [2mbfloat16[0m[4,256,768] │ │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_4/MlpBlock_0/Dense_0 │ Dense │ [2mbfloat16[0m[4,256,768] │ [2mbfloat16[0m[4,256,3072] │ bias: [2mfloat32[0m[3072] │
	│ │ │ │ │ kernel: [2mfloat32[0m[768,3072] │
	│ │ │ │ │ │
	│ │ │ │ │ [1m2,362,368 [0m[1;2m(9.4 MB)[0m │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_4/MlpBlock_0/Dropout_0 │ Dropout │ [2mbfloat16[0m[4,256,3072] │ [2mbfloat16[0m[4,256,3072] │ │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_4/MlpBlock_0/Dense_1 │ Dense │ [2mbfloat16[0m[4,256,3072] │ [2mbfloat16[0m[4,256,768] │ bias: [2mfloat32[0m[768] │
	│ │ │ │ │ kernel: [2mfloat32[0m[3072,768] │
	│ │ │ │ │ │
	│ │ │ │ │ [1m2,360,064 [0m[1;2m(9.4 MB)[0m │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_4/MlpBlock_0/Dropout_1 │ Dropout │ [2mbfloat16[0m[4,256,768] │ [2mbfloat16[0m[4,256,768] │ │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_5 │ DiTBlock │ - [2mbfloat16[0m[4,256,768] │ [2mbfloat16[0m[4,256,768] │ │
	│ │ │ - [2mfloat32[0m[1,768] │ │ │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_5/Dense_0 │ Dense │ [2mfloat32[0m[1,768] │ [2mbfloat16[0m[1,4608] │ bias: [2mfloat32[0m[4608] │
	│ │ │ │ │ kernel: [2mfloat32[0m[768,4608] │
	│ │ │ │ │ │
	│ │ │ │ │ [1m3,543,552 [0m[1;2m(14.2 MB)[0m │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_5/LayerNorm_0 │ LayerNorm │ [2mbfloat16[0m[4,256,768] │ [2mbfloat16[0m[4,256,768] │ │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_5/Dense_1 │ Dense │ [2mbfloat16[0m[4,256,768] │ [2mbfloat16[0m[4,256,768] │ bias: [2mfloat32[0m[768] │
	│ │ │ │ │ kernel: [2mfloat32[0m[768,768] │
	│ │ │ │ │ │
	│ │ │ │ │ [1m590,592 [0m[1;2m(2.4 MB)[0m │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_5/Dense_2 │ Dense │ [2mbfloat16[0m[4,256,768] │ [2mbfloat16[0m[4,256,768] │ bias: [2mfloat32[0m[768] │
	│ │ │ │ │ kernel: [2mfloat32[0m[768,768] │
	│ │ │ │ │ │
	│ │ │ │ │ [1m590,592 [0m[1;2m(2.4 MB)[0m │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_5/Dense_3 │ Dense │ [2mbfloat16[0m[4,256,768] │ [2mbfloat16[0m[4,256,768] │ bias: [2mfloat32[0m[768] │
	│ │ │ │ │ kernel: [2mfloat32[0m[768,768] │
	│ │ │ │ │ │
	│ │ │ │ │ [1m590,592 [0m[1;2m(2.4 MB)[0m │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_5/Dense_4 │ Dense │ [2mfloat32[0m[4,256,768] │ [2mbfloat16[0m[4,256,768] │ bias: [2mfloat32[0m[768] │
	│ │ │ │ │ kernel: [2mfloat32[0m[768,768] │
	│ │ │ │ │ │
	│ │ │ │ │ [1m590,592 [0m[1;2m(2.4 MB)[0m │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_5/LayerNorm_1 │ LayerNorm │ [2mbfloat16[0m[4,256,768] │ [2mbfloat16[0m[4,256,768] │ │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_5/MlpBlock_0 │ MlpBlock │ [2mbfloat16[0m[4,256,768] │ [2mbfloat16[0m[4,256,768] │ │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_5/MlpBlock_0/Dense_0 │ Dense │ [2mbfloat16[0m[4,256,768] │ [2mbfloat16[0m[4,256,3072] │ bias: [2mfloat32[0m[3072] │
	│ │ │ │ │ kernel: [2mfloat32[0m[768,3072] │
	│ │ │ │ │ │
	│ │ │ │ │ [1m2,362,368 [0m[1;2m(9.4 MB)[0m │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_5/MlpBlock_0/Dropout_0 │ Dropout │ [2mbfloat16[0m[4,256,3072] │ [2mbfloat16[0m[4,256,3072] │ │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_5/MlpBlock_0/Dense_1 │ Dense │ [2mbfloat16[0m[4,256,3072] │ [2mbfloat16[0m[4,256,768] │ bias: [2mfloat32[0m[768] │
	│ │ │ │ │ kernel: [2mfloat32[0m[3072,768] │
	│ │ │ │ │ │
	│ │ │ │ │ [1m2,360,064 [0m[1;2m(9.4 MB)[0m │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_5/MlpBlock_0/Dropout_1 │ Dropout │ [2mbfloat16[0m[4,256,768] │ [2mbfloat16[0m[4,256,768] │ │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_6 │ DiTBlock │ - [2mbfloat16[0m[4,256,768] │ [2mbfloat16[0m[4,256,768] │ │
	│ │ │ - [2mfloat32[0m[1,768] │ │ │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_6/Dense_0 │ Dense │ [2mfloat32[0m[1,768] │ [2mbfloat16[0m[1,4608] │ bias: [2mfloat32[0m[4608] │
	│ │ │ │ │ kernel: [2mfloat32[0m[768,4608] │
	│ │ │ │ │ │
	│ │ │ │ │ [1m3,543,552 [0m[1;2m(14.2 MB)[0m │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_6/LayerNorm_0 │ LayerNorm │ [2mbfloat16[0m[4,256,768] │ [2mbfloat16[0m[4,256,768] │ │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_6/Dense_1 │ Dense │ [2mbfloat16[0m[4,256,768] │ [2mbfloat16[0m[4,256,768] │ bias: [2mfloat32[0m[768] │
	│ │ │ │ │ kernel: [2mfloat32[0m[768,768] │
	│ │ │ │ │ │
	│ │ │ │ │ [1m590,592 [0m[1;2m(2.4 MB)[0m │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_6/Dense_2 │ Dense │ [2mbfloat16[0m[4,256,768] │ [2mbfloat16[0m[4,256,768] │ bias: [2mfloat32[0m[768] │
	│ │ │ │ │ kernel: [2mfloat32[0m[768,768] │
	│ │ │ │ │ │
	│ │ │ │ │ [1m590,592 [0m[1;2m(2.4 MB)[0m │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_6/Dense_3 │ Dense │ [2mbfloat16[0m[4,256,768] │ [2mbfloat16[0m[4,256,768] │ bias: [2mfloat32[0m[768] │
	│ │ │ │ │ kernel: [2mfloat32[0m[768,768] │
	│ │ │ │ │ │
	│ │ │ │ │ [1m590,592 [0m[1;2m(2.4 MB)[0m │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_6/Dense_4 │ Dense │ [2mfloat32[0m[4,256,768] │ [2mbfloat16[0m[4,256,768] │ bias: [2mfloat32[0m[768] │
	│ │ │ │ │ kernel: [2mfloat32[0m[768,768] │
	│ │ │ │ │ │
	│ │ │ │ │ [1m590,592 [0m[1;2m(2.4 MB)[0m │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_6/LayerNorm_1 │ LayerNorm │ [2mbfloat16[0m[4,256,768] │ [2mbfloat16[0m[4,256,768] │ │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_6/MlpBlock_0 │ MlpBlock │ [2mbfloat16[0m[4,256,768] │ [2mbfloat16[0m[4,256,768] │ │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_6/MlpBlock_0/Dense_0 │ Dense │ [2mbfloat16[0m[4,256,768] │ [2mbfloat16[0m[4,256,3072] │ bias: [2mfloat32[0m[3072] │
	│ │ │ │ │ kernel: [2mfloat32[0m[768,3072] │
	│ │ │ │ │ │
	│ │ │ │ │ [1m2,362,368 [0m[1;2m(9.4 MB)[0m │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_6/MlpBlock_0/Dropout_0 │ Dropout │ [2mbfloat16[0m[4,256,3072] │ [2mbfloat16[0m[4,256,3072] │ │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_6/MlpBlock_0/Dense_1 │ Dense │ [2mbfloat16[0m[4,256,3072] │ [2mbfloat16[0m[4,256,768] │ bias: [2mfloat32[0m[768] │
	│ │ │ │ │ kernel: [2mfloat32[0m[3072,768] │
	│ │ │ │ │ │
	│ │ │ │ │ [1m2,360,064 [0m[1;2m(9.4 MB)[0m │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_6/MlpBlock_0/Dropout_1 │ Dropout │ [2mbfloat16[0m[4,256,768] │ [2mbfloat16[0m[4,256,768] │ │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_7 │ DiTBlock │ - [2mbfloat16[0m[4,256,768] │ [2mbfloat16[0m[4,256,768] │ │
	│ │ │ - [2mfloat32[0m[1,768] │ │ │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_7/Dense_0 │ Dense │ [2mfloat32[0m[1,768] │ [2mbfloat16[0m[1,4608] │ bias: [2mfloat32[0m[4608] │
	│ │ │ │ │ kernel: [2mfloat32[0m[768,4608] │
	│ │ │ │ │ │
	│ │ │ │ │ [1m3,543,552 [0m[1;2m(14.2 MB)[0m │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_7/LayerNorm_0 │ LayerNorm │ [2mbfloat16[0m[4,256,768] │ [2mbfloat16[0m[4,256,768] │ │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_7/Dense_1 │ Dense │ [2mbfloat16[0m[4,256,768] │ [2mbfloat16[0m[4,256,768] │ bias: [2mfloat32[0m[768] │
	│ │ │ │ │ kernel: [2mfloat32[0m[768,768] │
	│ │ │ │ │ │
	│ │ │ │ │ [1m590,592 [0m[1;2m(2.4 MB)[0m │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_7/Dense_2 │ Dense │ [2mbfloat16[0m[4,256,768] │ [2mbfloat16[0m[4,256,768] │ bias: [2mfloat32[0m[768] │
	│ │ │ │ │ kernel: [2mfloat32[0m[768,768] │
	│ │ │ │ │ │
	│ │ │ │ │ [1m590,592 [0m[1;2m(2.4 MB)[0m │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_7/Dense_3 │ Dense │ [2mbfloat16[0m[4,256,768] │ [2mbfloat16[0m[4,256,768] │ bias: [2mfloat32[0m[768] │
	│ │ │ │ │ kernel: [2mfloat32[0m[768,768] │
	│ │ │ │ │ │
	│ │ │ │ │ [1m590,592 [0m[1;2m(2.4 MB)[0m │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_7/Dense_4 │ Dense │ [2mfloat32[0m[4,256,768] │ [2mbfloat16[0m[4,256,768] │ bias: [2mfloat32[0m[768] │
	│ │ │ │ │ kernel: [2mfloat32[0m[768,768] │
	│ │ │ │ │ │
	│ │ │ │ │ [1m590,592 [0m[1;2m(2.4 MB)[0m │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_7/LayerNorm_1 │ LayerNorm │ [2mbfloat16[0m[4,256,768] │ [2mbfloat16[0m[4,256,768] │ │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_7/MlpBlock_0 │ MlpBlock │ [2mbfloat16[0m[4,256,768] │ [2mbfloat16[0m[4,256,768] │ │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_7/MlpBlock_0/Dense_0 │ Dense │ [2mbfloat16[0m[4,256,768] │ [2mbfloat16[0m[4,256,3072] │ bias: [2mfloat32[0m[3072] │
	│ │ │ │ │ kernel: [2mfloat32[0m[768,3072] │
	│ │ │ │ │ │
	│ │ │ │ │ [1m2,362,368 [0m[1;2m(9.4 MB)[0m │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_7/MlpBlock_0/Dropout_0 │ Dropout │ [2mbfloat16[0m[4,256,3072] │ [2mbfloat16[0m[4,256,3072] │ │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_7/MlpBlock_0/Dense_1 │ Dense │ [2mbfloat16[0m[4,256,3072] │ [2mbfloat16[0m[4,256,768] │ bias: [2mfloat32[0m[768] │
	│ │ │ │ │ kernel: [2mfloat32[0m[3072,768] │
	│ │ │ │ │ │
	│ │ │ │ │ [1m2,360,064 [0m[1;2m(9.4 MB)[0m │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_7/MlpBlock_0/Dropout_1 │ Dropout │ [2mbfloat16[0m[4,256,768] │ [2mbfloat16[0m[4,256,768] │ │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_8 │ DiTBlock │ - [2mbfloat16[0m[4,256,768] │ [2mbfloat16[0m[4,256,768] │ │
	│ │ │ - [2mfloat32[0m[1,768] │ │ │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_8/Dense_0 │ Dense │ [2mfloat32[0m[1,768] │ [2mbfloat16[0m[1,4608] │ bias: [2mfloat32[0m[4608] │
	│ │ │ │ │ kernel: [2mfloat32[0m[768,4608] │
	│ │ │ │ │ │
	│ │ │ │ │ [1m3,543,552 [0m[1;2m(14.2 MB)[0m │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_8/LayerNorm_0 │ LayerNorm │ [2mbfloat16[0m[4,256,768] │ [2mbfloat16[0m[4,256,768] │ │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_8/Dense_1 │ Dense │ [2mbfloat16[0m[4,256,768] │ [2mbfloat16[0m[4,256,768] │ bias: [2mfloat32[0m[768] │
	│ │ │ │ │ kernel: [2mfloat32[0m[768,768] │
	│ │ │ │ │ │
	│ │ │ │ │ [1m590,592 [0m[1;2m(2.4 MB)[0m │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_8/Dense_2 │ Dense │ [2mbfloat16[0m[4,256,768] │ [2mbfloat16[0m[4,256,768] │ bias: [2mfloat32[0m[768] │
	│ │ │ │ │ kernel: [2mfloat32[0m[768,768] │
	│ │ │ │ │ │
	│ │ │ │ │ [1m590,592 [0m[1;2m(2.4 MB)[0m │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_8/Dense_3 │ Dense │ [2mbfloat16[0m[4,256,768] │ [2mbfloat16[0m[4,256,768] │ bias: [2mfloat32[0m[768] │
	│ │ │ │ │ kernel: [2mfloat32[0m[768,768] │
	│ │ │ │ │ │
	│ │ │ │ │ [1m590,592 [0m[1;2m(2.4 MB)[0m │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_8/Dense_4 │ Dense │ [2mfloat32[0m[4,256,768] │ [2mbfloat16[0m[4,256,768] │ bias: [2mfloat32[0m[768] │
	│ │ │ │ │ kernel: [2mfloat32[0m[768,768] │
	│ │ │ │ │ │
	│ │ │ │ │ [1m590,592 [0m[1;2m(2.4 MB)[0m │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_8/LayerNorm_1 │ LayerNorm │ [2mbfloat16[0m[4,256,768] │ [2mbfloat16[0m[4,256,768] │ │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_8/MlpBlock_0 │ MlpBlock │ [2mbfloat16[0m[4,256,768] │ [2mbfloat16[0m[4,256,768] │ │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_8/MlpBlock_0/Dense_0 │ Dense │ [2mbfloat16[0m[4,256,768] │ [2mbfloat16[0m[4,256,3072] │ bias: [2mfloat32[0m[3072] │
	│ │ │ │ │ kernel: [2mfloat32[0m[768,3072] │
	│ │ │ │ │ │
	│ │ │ │ │ [1m2,362,368 [0m[1;2m(9.4 MB)[0m │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_8/MlpBlock_0/Dropout_0 │ Dropout │ [2mbfloat16[0m[4,256,3072] │ [2mbfloat16[0m[4,256,3072] │ │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_8/MlpBlock_0/Dense_1 │ Dense │ [2mbfloat16[0m[4,256,3072] │ [2mbfloat16[0m[4,256,768] │ bias: [2mfloat32[0m[768] │
	│ │ │ │ │ kernel: [2mfloat32[0m[3072,768] │
	│ │ │ │ │ │
	│ │ │ │ │ [1m2,360,064 [0m[1;2m(9.4 MB)[0m │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_8/MlpBlock_0/Dropout_1 │ Dropout │ [2mbfloat16[0m[4,256,768] │ [2mbfloat16[0m[4,256,768] │ │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_9 │ DiTBlock │ - [2mbfloat16[0m[4,256,768] │ [2mbfloat16[0m[4,256,768] │ │
	│ │ │ - [2mfloat32[0m[1,768] │ │ │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_9/Dense_0 │ Dense │ [2mfloat32[0m[1,768] │ [2mbfloat16[0m[1,4608] │ bias: [2mfloat32[0m[4608] │
	│ │ │ │ │ kernel: [2mfloat32[0m[768,4608] │
	│ │ │ │ │ │
	│ │ │ │ │ [1m3,543,552 [0m[1;2m(14.2 MB)[0m │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_9/LayerNorm_0 │ LayerNorm │ [2mbfloat16[0m[4,256,768] │ [2mbfloat16[0m[4,256,768] │ │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_9/Dense_1 │ Dense │ [2mbfloat16[0m[4,256,768] │ [2mbfloat16[0m[4,256,768] │ bias: [2mfloat32[0m[768] │
	│ │ │ │ │ kernel: [2mfloat32[0m[768,768] │
	│ │ │ │ │ │
	│ │ │ │ │ [1m590,592 [0m[1;2m(2.4 MB)[0m │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_9/Dense_2 │ Dense │ [2mbfloat16[0m[4,256,768] │ [2mbfloat16[0m[4,256,768] │ bias: [2mfloat32[0m[768] │
	│ │ │ │ │ kernel: [2mfloat32[0m[768,768] │
	│ │ │ │ │ │
	│ │ │ │ │ [1m590,592 [0m[1;2m(2.4 MB)[0m │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_9/Dense_3 │ Dense │ [2mbfloat16[0m[4,256,768] │ [2mbfloat16[0m[4,256,768] │ bias: [2mfloat32[0m[768] │
	│ │ │ │ │ kernel: [2mfloat32[0m[768,768] │
	│ │ │ │ │ │
	│ │ │ │ │ [1m590,592 [0m[1;2m(2.4 MB)[0m │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_9/Dense_4 │ Dense │ [2mfloat32[0m[4,256,768] │ [2mbfloat16[0m[4,256,768] │ bias: [2mfloat32[0m[768] │
	│ │ │ │ │ kernel: [2mfloat32[0m[768,768] │
	│ │ │ │ │ │
	│ │ │ │ │ [1m590,592 [0m[1;2m(2.4 MB)[0m │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_9/LayerNorm_1 │ LayerNorm │ [2mbfloat16[0m[4,256,768] │ [2mbfloat16[0m[4,256,768] │ │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_9/MlpBlock_0 │ MlpBlock │ [2mbfloat16[0m[4,256,768] │ [2mbfloat16[0m[4,256,768] │ │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_9/MlpBlock_0/Dense_0 │ Dense │ [2mbfloat16[0m[4,256,768] │ [2mbfloat16[0m[4,256,3072] │ bias: [2mfloat32[0m[3072] │
	│ │ │ │ │ kernel: [2mfloat32[0m[768,3072] │
	│ │ │ │ │ │
	│ │ │ │ │ [1m2,362,368 [0m[1;2m(9.4 MB)[0m │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_9/MlpBlock_0/Dropout_0 │ Dropout │ [2mbfloat16[0m[4,256,3072] │ [2mbfloat16[0m[4,256,3072] │ │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_9/MlpBlock_0/Dense_1 │ Dense │ [2mbfloat16[0m[4,256,3072] │ [2mbfloat16[0m[4,256,768] │ bias: [2mfloat32[0m[768] │
	│ │ │ │ │ kernel: [2mfloat32[0m[3072,768] │
	│ │ │ │ │ │
	│ │ │ │ │ [1m2,360,064 [0m[1;2m(9.4 MB)[0m │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_9/MlpBlock_0/Dropout_1 │ Dropout │ [2mbfloat16[0m[4,256,768] │ [2mbfloat16[0m[4,256,768] │ │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_10 │ DiTBlock │ - [2mbfloat16[0m[4,256,768] │ [2mbfloat16[0m[4,256,768] │ │
	│ │ │ - [2mfloat32[0m[1,768] │ │ │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_10/Dense_0 │ Dense │ [2mfloat32[0m[1,768] │ [2mbfloat16[0m[1,4608] │ bias: [2mfloat32[0m[4608] │
	│ │ │ │ │ kernel: [2mfloat32[0m[768,4608] │
	│ │ │ │ │ │
	│ │ │ │ │ [1m3,543,552 [0m[1;2m(14.2 MB)[0m │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_10/LayerNorm_0 │ LayerNorm │ [2mbfloat16[0m[4,256,768] │ [2mbfloat16[0m[4,256,768] │ │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_10/Dense_1 │ Dense │ [2mbfloat16[0m[4,256,768] │ [2mbfloat16[0m[4,256,768] │ bias: [2mfloat32[0m[768] │
	│ │ │ │ │ kernel: [2mfloat32[0m[768,768] │
	│ │ │ │ │ │
	│ │ │ │ │ [1m590,592 [0m[1;2m(2.4 MB)[0m │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_10/Dense_2 │ Dense │ [2mbfloat16[0m[4,256,768] │ [2mbfloat16[0m[4,256,768] │ bias: [2mfloat32[0m[768] │
	│ │ │ │ │ kernel: [2mfloat32[0m[768,768] │
	│ │ │ │ │ │
	│ │ │ │ │ [1m590,592 [0m[1;2m(2.4 MB)[0m │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_10/Dense_3 │ Dense │ [2mbfloat16[0m[4,256,768] │ [2mbfloat16[0m[4,256,768] │ bias: [2mfloat32[0m[768] │
	│ │ │ │ │ kernel: [2mfloat32[0m[768,768] │
	│ │ │ │ │ │
	│ │ │ │ │ [1m590,592 [0m[1;2m(2.4 MB)[0m │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_10/Dense_4 │ Dense │ [2mfloat32[0m[4,256,768] │ [2mbfloat16[0m[4,256,768] │ bias: [2mfloat32[0m[768] │
	│ │ │ │ │ kernel: [2mfloat32[0m[768,768] │
	│ │ │ │ │ │
	│ │ │ │ │ [1m590,592 [0m[1;2m(2.4 MB)[0m │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_10/LayerNorm_1 │ LayerNorm │ [2mbfloat16[0m[4,256,768] │ [2mbfloat16[0m[4,256,768] │ │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_10/MlpBlock_0 │ MlpBlock │ [2mbfloat16[0m[4,256,768] │ [2mbfloat16[0m[4,256,768] │ │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_10/MlpBlock_0/Dense_0 │ Dense │ [2mbfloat16[0m[4,256,768] │ [2mbfloat16[0m[4,256,3072] │ bias: [2mfloat32[0m[3072] │
	│ │ │ │ │ kernel: [2mfloat32[0m[768,3072] │
	│ │ │ │ │ │
	│ │ │ │ │ [1m2,362,368 [0m[1;2m(9.4 MB)[0m │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_10/MlpBlock_0/Dropout_0 │ Dropout │ [2mbfloat16[0m[4,256,3072] │ [2mbfloat16[0m[4,256,3072] │ │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_10/MlpBlock_0/Dense_1 │ Dense │ [2mbfloat16[0m[4,256,3072] │ [2mbfloat16[0m[4,256,768] │ bias: [2mfloat32[0m[768] │
	│ │ │ │ │ kernel: [2mfloat32[0m[3072,768] │
	│ │ │ │ │ │
	│ │ │ │ │ [1m2,360,064 [0m[1;2m(9.4 MB)[0m │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_10/MlpBlock_0/Dropout_1 │ Dropout │ [2mbfloat16[0m[4,256,768] │ [2mbfloat16[0m[4,256,768] │ │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_11 │ DiTBlock │ - [2mbfloat16[0m[4,256,768] │ [2mbfloat16[0m[4,256,768] │ │
	│ │ │ - [2mfloat32[0m[1,768] │ │ │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_11/Dense_0 │ Dense │ [2mfloat32[0m[1,768] │ [2mbfloat16[0m[1,4608] │ bias: [2mfloat32[0m[4608] │
	│ │ │ │ │ kernel: [2mfloat32[0m[768,4608] │
	│ │ │ │ │ │
	│ │ │ │ │ [1m3,543,552 [0m[1;2m(14.2 MB)[0m │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_11/LayerNorm_0 │ LayerNorm │ [2mbfloat16[0m[4,256,768] │ [2mbfloat16[0m[4,256,768] │ │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_11/Dense_1 │ Dense │ [2mbfloat16[0m[4,256,768] │ [2mbfloat16[0m[4,256,768] │ bias: [2mfloat32[0m[768] │
	│ │ │ │ │ kernel: [2mfloat32[0m[768,768] │
	│ │ │ │ │ │
	│ │ │ │ │ [1m590,592 [0m[1;2m(2.4 MB)[0m │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_11/Dense_2 │ Dense │ [2mbfloat16[0m[4,256,768] │ [2mbfloat16[0m[4,256,768] │ bias: [2mfloat32[0m[768] │
	│ │ │ │ │ kernel: [2mfloat32[0m[768,768] │
	│ │ │ │ │ │
	│ │ │ │ │ [1m590,592 [0m[1;2m(2.4 MB)[0m │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_11/Dense_3 │ Dense │ [2mbfloat16[0m[4,256,768] │ [2mbfloat16[0m[4,256,768] │ bias: [2mfloat32[0m[768] │
	│ │ │ │ │ kernel: [2mfloat32[0m[768,768] │
	│ │ │ │ │ │
	│ │ │ │ │ [1m590,592 [0m[1;2m(2.4 MB)[0m │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_11/Dense_4 │ Dense │ [2mfloat32[0m[4,256,768] │ [2mbfloat16[0m[4,256,768] │ bias: [2mfloat32[0m[768] │
	│ │ │ │ │ kernel: [2mfloat32[0m[768,768] │
	│ │ │ │ │ │
	│ │ │ │ │ [1m590,592 [0m[1;2m(2.4 MB)[0m │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_11/LayerNorm_1 │ LayerNorm │ [2mbfloat16[0m[4,256,768] │ [2mbfloat16[0m[4,256,768] │ │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_11/MlpBlock_0 │ MlpBlock │ [2mbfloat16[0m[4,256,768] │ [2mbfloat16[0m[4,256,768] │ │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_11/MlpBlock_0/Dense_0 │ Dense │ [2mbfloat16[0m[4,256,768] │ [2mbfloat16[0m[4,256,3072] │ bias: [2mfloat32[0m[3072] │
	│ │ │ │ │ kernel: [2mfloat32[0m[768,3072] │
	│ │ │ │ │ │
	│ │ │ │ │ [1m2,362,368 [0m[1;2m(9.4 MB)[0m │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_11/MlpBlock_0/Dropout_0 │ Dropout │ [2mbfloat16[0m[4,256,3072] │ [2mbfloat16[0m[4,256,3072] │ │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_11/MlpBlock_0/Dense_1 │ Dense │ [2mbfloat16[0m[4,256,3072] │ [2mbfloat16[0m[4,256,768] │ bias: [2mfloat32[0m[768] │
	│ │ │ │ │ kernel: [2mfloat32[0m[3072,768] │
	│ │ │ │ │ │
	│ │ │ │ │ [1m2,360,064 [0m[1;2m(9.4 MB)[0m │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ DiTBlock_11/MlpBlock_0/Dropout_1 │ Dropout │ [2mbfloat16[0m[4,256,768] │ [2mbfloat16[0m[4,256,768] │ │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ FinalLayer_0 │ FinalLayer │ - [2mbfloat16[0m[4,256,768] │ [2mbfloat16[0m[4,256,16] │ │
	│ │ │ - [2mfloat32[0m[1,768] │ │ │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ FinalLayer_0/Dense_0 │ Dense │ [2mfloat32[0m[1,768] │ [2mbfloat16[0m[1,1536] │ bias: [2mfloat32[0m[1536] │
	│ │ │ │ │ kernel: [2mfloat32[0m[768,1536] │
	│ │ │ │ │ │
	│ │ │ │ │ [1m1,181,184 [0m[1;2m(4.7 MB)[0m │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ FinalLayer_0/LayerNorm_0 │ LayerNorm │ [2mbfloat16[0m[4,256,768] │ [2mbfloat16[0m[4,256,768] │ │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ FinalLayer_0/Dense_1 │ Dense │ [2mbfloat16[0m[4,256,768] │ [2mbfloat16[0m[4,256,16] │ bias: [2mfloat32[0m[16] │
	│ │ │ │ │ kernel: [2mfloat32[0m[768,16] │
	│ │ │ │ │ │
	│ │ │ │ │ [1m12,304 [0m[1;2m(49.2 KB)[0m │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│ Embed_0 │ Embed │ [2mint32[0m[1] │ [2mfloat32[0m[1,1] │ embedding: [2mfloat32[0m[256,1] │
	│ │ │ │ │ │
	│ │ │ │ │ [1m256 [0m[1;2m(1.0 KB)[0m │
	├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────────┼──────────────────────────────┤
	│[1m [0m[1m [0m[1m [0m│[1m [0m[1m [0m[1m [0m│[1m [0m[1m [0m[1m [0m│[1m [0m[1m Total[0m[1m [0m│[1m [0m[1m131,091,728 [0m[1;2m(524.4 MB)[0m[1m [0m[1m [0m│
	└──────────────────────────────────┴──────────────────┴───────────────────────┴───────────────────────┴──────────────────────────────┘
	[1m [0m
	[1m Total Parameters: 131,091,728 [0m[1;2m(524.4 MB)[0m[1m [0m


	DiT: Input of shape (4, 32, 32, 4) dtype float32
	DiT: After patch embed, shape is (4, 256, 768) dtype bfloat16
	DiT: Patch Embed of shape (4, 256, 768) dtype bfloat16
	DiT: Conditioning of shape (1, 768) dtype float32
	Loaded checkpoint from 6459 seconds ago.

	parameter shapes:
	('PatchEmbed_0', 'Conv_0', 'kernel'): (2, 2, 4, 768)
	('PatchEmbed_0', 'Conv_0', 'bias'): (768,)
	('TimestepEmbedder_0', 'Dense_0', 'kernel'): (256, 768)
	('TimestepEmbedder_0', 'Dense_0', 'bias'): (768,)
	('TimestepEmbedder_0', 'Dense_1', 'kernel'): (768, 768)
	('TimestepEmbedder_0', 'Dense_1', 'bias'): (768,)
	('TimestepEmbedder_1', 'Dense_0', 'kernel'): (256, 768)
	('TimestepEmbedder_1', 'Dense_0', 'bias'): (768,)
	('TimestepEmbedder_1', 'Dense_1', 'kernel'): (768, 768)
	('TimestepEmbedder_1', 'Dense_1', 'bias'): (768,)
	('LabelEmbedder_0', 'Embed_0', 'embedding'): (1001, 768)
	('DiTBlock_0', 'Dense_0', 'kernel'): (768, 4608)
	('DiTBlock_0', 'Dense_0', 'bias'): (4608,)
	('DiTBlock_0', 'Dense_1', 'kernel'): (768, 768)
	('DiTBlock_0', 'Dense_1', 'bias'): (768,)
	('DiTBlock_0', 'Dense_2', 'kernel'): (768, 768)
	('DiTBlock_0', 'Dense_2', 'bias'): (768,)
	('DiTBlock_0', 'Dense_3', 'kernel'): (768, 768)
	('DiTBlock_0', 'Dense_3', 'bias'): (768,)
	('DiTBlock_0', 'Dense_4', 'kernel'): (768, 768)
	('DiTBlock_0', 'Dense_4', 'bias'): (768,)
	('DiTBlock_0', 'MlpBlock_0', 'Dense_0', 'kernel'): (768, 3072)
	('DiTBlock_0', 'MlpBlock_0', 'Dense_0', 'bias'): (3072,)
	('DiTBlock_0', 'MlpBlock_0', 'Dense_1', 'kernel'): (3072, 768)
	('DiTBlock_0', 'MlpBlock_0', 'Dense_1', 'bias'): (768,)
	('DiTBlock_1', 'Dense_0', 'kernel'): (768, 4608)
	('DiTBlock_1', 'Dense_0', 'bias'): (4608,)
	('DiTBlock_1', 'Dense_1', 'kernel'): (768, 768)
	('DiTBlock_1', 'Dense_1', 'bias'): (768,)
	('DiTBlock_1', 'Dense_2', 'kernel'): (768, 768)
	('DiTBlock_1', 'Dense_2', 'bias'): (768,)
	('DiTBlock_1', 'Dense_3', 'kernel'): (768, 768)
	('DiTBlock_1', 'Dense_3', 'bias'): (768,)
	('DiTBlock_1', 'Dense_4', 'kernel'): (768, 768)
	('DiTBlock_1', 'Dense_4', 'bias'): (768,)
	('DiTBlock_1', 'MlpBlock_0', 'Dense_0', 'kernel'): (768, 3072)
	('DiTBlock_1', 'MlpBlock_0', 'Dense_0', 'bias'): (3072,)
	('DiTBlock_1', 'MlpBlock_0', 'Dense_1', 'kernel'): (3072, 768)
	('DiTBlock_1', 'MlpBlock_0', 'Dense_1', 'bias'): (768,)
	('DiTBlock_2', 'Dense_0', 'kernel'): (768, 4608)
	('DiTBlock_2', 'Dense_0', 'bias'): (4608,)
	('DiTBlock_2', 'Dense_1', 'kernel'): (768, 768)
	('DiTBlock_2', 'Dense_1', 'bias'): (768,)
	('DiTBlock_2', 'Dense_2', 'kernel'): (768, 768)
	('DiTBlock_2', 'Dense_2', 'bias'): (768,)
	('DiTBlock_2', 'Dense_3', 'kernel'): (768, 768)
	('DiTBlock_2', 'Dense_3', 'bias'): (768,)
	('DiTBlock_2', 'Dense_4', 'kernel'): (768, 768)
	('DiTBlock_2', 'Dense_4', 'bias'): (768,)
	('DiTBlock_2', 'MlpBlock_0', 'Dense_0', 'kernel'): (768, 3072)
	('DiTBlock_2', 'MlpBlock_0', 'Dense_0', 'bias'): (3072,)
	('DiTBlock_2', 'MlpBlock_0', 'Dense_1', 'kernel'): (3072, 768)
	('DiTBlock_2', 'MlpBlock_0', 'Dense_1', 'bias'): (768,)
	('DiTBlock_3', 'Dense_0', 'kernel'): (768, 4608)
	('DiTBlock_3', 'Dense_0', 'bias'): (4608,)
	('DiTBlock_3', 'Dense_1', 'kernel'): (768, 768)
	('DiTBlock_3', 'Dense_1', 'bias'): (768,)
	('DiTBlock_3', 'Dense_2', 'kernel'): (768, 768)
	('DiTBlock_3', 'Dense_2', 'bias'): (768,)
	('DiTBlock_3', 'Dense_3', 'kernel'): (768, 768)
	('DiTBlock_3', 'Dense_3', 'bias'): (768,)
	('DiTBlock_3', 'Dense_4', 'kernel'): (768, 768)
	('DiTBlock_3', 'Dense_4', 'bias'): (768,)
	('DiTBlock_3', 'MlpBlock_0', 'Dense_0', 'kernel'): (768, 3072)
	('DiTBlock_3', 'MlpBlock_0', 'Dense_0', 'bias'): (3072,)
	('DiTBlock_3', 'MlpBlock_0', 'Dense_1', 'kernel'): (3072, 768)
	('DiTBlock_3', 'MlpBlock_0', 'Dense_1', 'bias'): (768,)
	('DiTBlock_4', 'Dense_0', 'kernel'): (768, 4608)
	('DiTBlock_4', 'Dense_0', 'bias'): (4608,)
	('DiTBlock_4', 'Dense_1', 'kernel'): (768, 768)
	('DiTBlock_4', 'Dense_1', 'bias'): (768,)
	('DiTBlock_4', 'Dense_2', 'kernel'): (768, 768)
	('DiTBlock_4', 'Dense_2', 'bias'): (768,)
	('DiTBlock_4', 'Dense_3', 'kernel'): (768, 768)
	('DiTBlock_4', 'Dense_3', 'bias'): (768,)
	('DiTBlock_4', 'Dense_4', 'kernel'): (768, 768)
	('DiTBlock_4', 'Dense_4', 'bias'): (768,)
	('DiTBlock_4', 'MlpBlock_0', 'Dense_0', 'kernel'): (768, 3072)
	('DiTBlock_4', 'MlpBlock_0', 'Dense_0', 'bias'): (3072,)
	('DiTBlock_4', 'MlpBlock_0', 'Dense_1', 'kernel'): (3072, 768)
	('DiTBlock_4', 'MlpBlock_0', 'Dense_1', 'bias'): (768,)
	('DiTBlock_5', 'Dense_0', 'kernel'): (768, 4608)
	('DiTBlock_5', 'Dense_0', 'bias'): (4608,)
	('DiTBlock_5', 'Dense_1', 'kernel'): (768, 768)
	('DiTBlock_5', 'Dense_1', 'bias'): (768,)
	('DiTBlock_5', 'Dense_2', 'kernel'): (768, 768)
	('DiTBlock_5', 'Dense_2', 'bias'): (768,)
	('DiTBlock_5', 'Dense_3', 'kernel'): (768, 768)
	('DiTBlock_5', 'Dense_3', 'bias'): (768,)
	('DiTBlock_5', 'Dense_4', 'kernel'): (768, 768)
	('DiTBlock_5', 'Dense_4', 'bias'): (768,)
	('DiTBlock_5', 'MlpBlock_0', 'Dense_0', 'kernel'): (768, 3072)
	('DiTBlock_5', 'MlpBlock_0', 'Dense_0', 'bias'): (3072,)
	('DiTBlock_5', 'MlpBlock_0', 'Dense_1', 'kernel'): (3072, 768)
	('DiTBlock_5', 'MlpBlock_0', 'Dense_1', 'bias'): (768,)
	('DiTBlock_6', 'Dense_0', 'kernel'): (768, 4608)
	('DiTBlock_6', 'Dense_0', 'bias'): (4608,)
	('DiTBlock_6', 'Dense_1', 'kernel'): (768, 768)
	('DiTBlock_6', 'Dense_1', 'bias'): (768,)
	('DiTBlock_6', 'Dense_2', 'kernel'): (768, 768)
	('DiTBlock_6', 'Dense_2', 'bias'): (768,)
	('DiTBlock_6', 'Dense_3', 'kernel'): (768, 768)
	('DiTBlock_6', 'Dense_3', 'bias'): (768,)
	('DiTBlock_6', 'Dense_4', 'kernel'): (768, 768)
	('DiTBlock_6', 'Dense_4', 'bias'): (768,)
	('DiTBlock_6', 'MlpBlock_0', 'Dense_0', 'kernel'): (768, 3072)
	('DiTBlock_6', 'MlpBlock_0', 'Dense_0', 'bias'): (3072,)
	('DiTBlock_6', 'MlpBlock_0', 'Dense_1', 'kernel'): (3072, 768)
	('DiTBlock_6', 'MlpBlock_0', 'Dense_1', 'bias'): (768,)
	('DiTBlock_7', 'Dense_0', 'kernel'): (768, 4608)
	('DiTBlock_7', 'Dense_0', 'bias'): (4608,)
	('DiTBlock_7', 'Dense_1', 'kernel'): (768, 768)
	('DiTBlock_7', 'Dense_1', 'bias'): (768,)
	('DiTBlock_7', 'Dense_2', 'kernel'): (768, 768)
	('DiTBlock_7', 'Dense_2', 'bias'): (768,)
	('DiTBlock_7', 'Dense_3', 'kernel'): (768, 768)
	('DiTBlock_7', 'Dense_3', 'bias'): (768,)
	('DiTBlock_7', 'Dense_4', 'kernel'): (768, 768)
	('DiTBlock_7', 'Dense_4', 'bias'): (768,)
	('DiTBlock_7', 'MlpBlock_0', 'Dense_0', 'kernel'): (768, 3072)
	('DiTBlock_7', 'MlpBlock_0', 'Dense_0', 'bias'): (3072,)
	('DiTBlock_7', 'MlpBlock_0', 'Dense_1', 'kernel'): (3072, 768)
	('DiTBlock_7', 'MlpBlock_0', 'Dense_1', 'bias'): (768,)
	('DiTBlock_8', 'Dense_0', 'kernel'): (768, 4608)
	('DiTBlock_8', 'Dense_0', 'bias'): (4608,)
	('DiTBlock_8', 'Dense_1', 'kernel'): (768, 768)
	('DiTBlock_8', 'Dense_1', 'bias'): (768,)
	('DiTBlock_8', 'Dense_2', 'kernel'): (768, 768)
	('DiTBlock_8', 'Dense_2', 'bias'): (768,)
	('DiTBlock_8', 'Dense_3', 'kernel'): (768, 768)
	('DiTBlock_8', 'Dense_3', 'bias'): (768,)
	('DiTBlock_8', 'Dense_4', 'kernel'): (768, 768)
	('DiTBlock_8', 'Dense_4', 'bias'): (768,)
	('DiTBlock_8', 'MlpBlock_0', 'Dense_0', 'kernel'): (768, 3072)
	('DiTBlock_8', 'MlpBlock_0', 'Dense_0', 'bias'): (3072,)
	('DiTBlock_8', 'MlpBlock_0', 'Dense_1', 'kernel'): (3072, 768)
	('DiTBlock_8', 'MlpBlock_0', 'Dense_1', 'bias'): (768,)
	('DiTBlock_9', 'Dense_0', 'kernel'): (768, 4608)
	('DiTBlock_9', 'Dense_0', 'bias'): (4608,)
	('DiTBlock_9', 'Dense_1', 'kernel'): (768, 768)
	('DiTBlock_9', 'Dense_1', 'bias'): (768,)
	('DiTBlock_9', 'Dense_2', 'kernel'): (768, 768)
	('DiTBlock_9', 'Dense_2', 'bias'): (768,)
	('DiTBlock_9', 'Dense_3', 'kernel'): (768, 768)
	('DiTBlock_9', 'Dense_3', 'bias'): (768,)
	('DiTBlock_9', 'Dense_4', 'kernel'): (768, 768)
	('DiTBlock_9', 'Dense_4', 'bias'): (768,)
	('DiTBlock_9', 'MlpBlock_0', 'Dense_0', 'kernel'): (768, 3072)
	('DiTBlock_9', 'MlpBlock_0', 'Dense_0', 'bias'): (3072,)
	('DiTBlock_9', 'MlpBlock_0', 'Dense_1', 'kernel'): (3072, 768)
	('DiTBlock_9', 'MlpBlock_0', 'Dense_1', 'bias'): (768,)
	('DiTBlock_10', 'Dense_0', 'kernel'): (768, 4608)
	('DiTBlock_10', 'Dense_0', 'bias'): (4608,)
	('DiTBlock_10', 'Dense_1', 'kernel'): (768, 768)
	('DiTBlock_10', 'Dense_1', 'bias'): (768,)
	('DiTBlock_10', 'Dense_2', 'kernel'): (768, 768)
	('DiTBlock_10', 'Dense_2', 'bias'): (768,)
	('DiTBlock_10', 'Dense_3', 'kernel'): (768, 768)
	('DiTBlock_10', 'Dense_3', 'bias'): (768,)
	('DiTBlock_10', 'Dense_4', 'kernel'): (768, 768)
	('DiTBlock_10', 'Dense_4', 'bias'): (768,)
	('DiTBlock_10', 'MlpBlock_0', 'Dense_0', 'kernel'): (768, 3072)
	('DiTBlock_10', 'MlpBlock_0', 'Dense_0', 'bias'): (3072,)
	('DiTBlock_10', 'MlpBlock_0', 'Dense_1', 'kernel'): (3072, 768)
	('DiTBlock_10', 'MlpBlock_0', 'Dense_1', 'bias'): (768,)
	('DiTBlock_11', 'Dense_0', 'kernel'): (768, 4608)
	('DiTBlock_11', 'Dense_0', 'bias'): (4608,)
	('DiTBlock_11', 'Dense_1', 'kernel'): (768, 768)
	('DiTBlock_11', 'Dense_1', 'bias'): (768,)
	('DiTBlock_11', 'Dense_2', 'kernel'): (768, 768)
	('DiTBlock_11', 'Dense_2', 'bias'): (768,)
	('DiTBlock_11', 'Dense_3', 'kernel'): (768, 768)
	('DiTBlock_11', 'Dense_3', 'bias'): (768,)
	('DiTBlock_11', 'Dense_4', 'kernel'): (768, 768)
	('DiTBlock_11', 'Dense_4', 'bias'): (768,)
	('DiTBlock_11', 'MlpBlock_0', 'Dense_0', 'kernel'): (768, 3072)
	('DiTBlock_11', 'MlpBlock_0', 'Dense_0', 'bias'): (3072,)
	('DiTBlock_11', 'MlpBlock_0', 'Dense_1', 'kernel'): (3072, 768)
	('DiTBlock_11', 'MlpBlock_0', 'Dense_1', 'bias'): (768,)
	('FinalLayer_0', 'Dense_0', 'kernel'): (768, 1536)
	('FinalLayer_0', 'Dense_0', 'bias'): (1536,)
	('FinalLayer_0', 'Dense_1', 'kernel'): (768, 16)
	('FinalLayer_0', 'Dense_1', 'bias'): (16,)
	('Embed_0', 'embedding'): (256, 1)

	parameter shapes:
	('DiTBlock_0', 'Dense_0', 'bias'): (1, 4608)
	('DiTBlock_0', 'Dense_0', 'kernel'): (1, 768, 4608)
	('DiTBlock_0', 'Dense_1', 'bias'): (1, 768)
	('DiTBlock_0', 'Dense_1', 'kernel'): (1, 768, 768)
	('DiTBlock_0', 'Dense_2', 'bias'): (1, 768)
	('DiTBlock_0', 'Dense_2', 'kernel'): (1, 768, 768)
	('DiTBlock_0', 'Dense_3', 'bias'): (1, 768)
	('DiTBlock_0', 'Dense_3', 'kernel'): (1, 768, 768)
	('DiTBlock_0', 'Dense_4', 'bias'): (1, 768)
	('DiTBlock_0', 'Dense_4', 'kernel'): (1, 768, 768)
	('DiTBlock_0', 'MlpBlock_0', 'Dense_0', 'bias'): (1, 3072)
	('DiTBlock_0', 'MlpBlock_0', 'Dense_0', 'kernel'): (1, 768, 3072)
	('DiTBlock_0', 'MlpBlock_0', 'Dense_1', 'bias'): (1, 768)
	('DiTBlock_0', 'MlpBlock_0', 'Dense_1', 'kernel'): (1, 3072, 768)
	('DiTBlock_1', 'Dense_0', 'bias'): (1, 4608)
	('DiTBlock_1', 'Dense_0', 'kernel'): (1, 768, 4608)
	('DiTBlock_1', 'Dense_1', 'bias'): (1, 768)
	('DiTBlock_1', 'Dense_1', 'kernel'): (1, 768, 768)
	('DiTBlock_1', 'Dense_2', 'bias'): (1, 768)
	('DiTBlock_1', 'Dense_2', 'kernel'): (1, 768, 768)
	('DiTBlock_1', 'Dense_3', 'bias'): (1, 768)
	('DiTBlock_1', 'Dense_3', 'kernel'): (1, 768, 768)
	('DiTBlock_1', 'Dense_4', 'bias'): (1, 768)
	('DiTBlock_1', 'Dense_4', 'kernel'): (1, 768, 768)
	('DiTBlock_1', 'MlpBlock_0', 'Dense_0', 'bias'): (1, 3072)
	('DiTBlock_1', 'MlpBlock_0', 'Dense_0', 'kernel'): (1, 768, 3072)
	('DiTBlock_1', 'MlpBlock_0', 'Dense_1', 'bias'): (1, 768)
	('DiTBlock_1', 'MlpBlock_0', 'Dense_1', 'kernel'): (1, 3072, 768)
	('DiTBlock_10', 'Dense_0', 'bias'): (1, 4608)
	('DiTBlock_10', 'Dense_0', 'kernel'): (1, 768, 4608)
	('DiTBlock_10', 'Dense_1', 'bias'): (1, 768)
	('DiTBlock_10', 'Dense_1', 'kernel'): (1, 768, 768)
	('DiTBlock_10', 'Dense_2', 'bias'): (1, 768)
	('DiTBlock_10', 'Dense_2', 'kernel'): (1, 768, 768)
	('DiTBlock_10', 'Dense_3', 'bias'): (1, 768)
	('DiTBlock_10', 'Dense_3', 'kernel'): (1, 768, 768)
	('DiTBlock_10', 'Dense_4', 'bias'): (1, 768)
	('DiTBlock_10', 'Dense_4', 'kernel'): (1, 768, 768)
	('DiTBlock_10', 'MlpBlock_0', 'Dense_0', 'bias'): (1, 3072)
	('DiTBlock_10', 'MlpBlock_0', 'Dense_0', 'kernel'): (1, 768, 3072)
	('DiTBlock_10', 'MlpBlock_0', 'Dense_1', 'bias'): (1, 768)
	('DiTBlock_10', 'MlpBlock_0', 'Dense_1', 'kernel'): (1, 3072, 768)
	('DiTBlock_11', 'Dense_0', 'bias'): (1, 4608)
	('DiTBlock_11', 'Dense_0', 'kernel'): (1, 768, 4608)
	('DiTBlock_11', 'Dense_1', 'bias'): (1, 768)
	('DiTBlock_11', 'Dense_1', 'kernel'): (1, 768, 768)
	('DiTBlock_11', 'Dense_2', 'bias'): (1, 768)
	('DiTBlock_11', 'Dense_2', 'kernel'): (1, 768, 768)
	('DiTBlock_11', 'Dense_3', 'bias'): (1, 768)
	('DiTBlock_11', 'Dense_3', 'kernel'): (1, 768, 768)
	('DiTBlock_11', 'Dense_4', 'bias'): (1, 768)
	('DiTBlock_11', 'Dense_4', 'kernel'): (1, 768, 768)
	('DiTBlock_11', 'MlpBlock_0', 'Dense_0', 'bias'): (1, 3072)
	('DiTBlock_11', 'MlpBlock_0', 'Dense_0', 'kernel'): (1, 768, 3072)
	('DiTBlock_11', 'MlpBlock_0', 'Dense_1', 'bias'): (1, 768)
	('DiTBlock_11', 'MlpBlock_0', 'Dense_1', 'kernel'): (1, 3072, 768)
	('DiTBlock_2', 'Dense_0', 'bias'): (1, 4608)
	('DiTBlock_2', 'Dense_0', 'kernel'): (1, 768, 4608)
	('DiTBlock_2', 'Dense_1', 'bias'): (1, 768)
	('DiTBlock_2', 'Dense_1', 'kernel'): (1, 768, 768)
	('DiTBlock_2', 'Dense_2', 'bias'): (1, 768)
	('DiTBlock_2', 'Dense_2', 'kernel'): (1, 768, 768)
	('DiTBlock_2', 'Dense_3', 'bias'): (1, 768)
	('DiTBlock_2', 'Dense_3', 'kernel'): (1, 768, 768)
	('DiTBlock_2', 'Dense_4', 'bias'): (1, 768)
	('DiTBlock_2', 'Dense_4', 'kernel'): (1, 768, 768)
	('DiTBlock_2', 'MlpBlock_0', 'Dense_0', 'bias'): (1, 3072)
	('DiTBlock_2', 'MlpBlock_0', 'Dense_0', 'kernel'): (1, 768, 3072)
	('DiTBlock_2', 'MlpBlock_0', 'Dense_1', 'bias'): (1, 768)
	('DiTBlock_2', 'MlpBlock_0', 'Dense_1', 'kernel'): (1, 3072, 768)
	('DiTBlock_3', 'Dense_0', 'bias'): (1, 4608)
	('DiTBlock_3', 'Dense_0', 'kernel'): (1, 768, 4608)
	('DiTBlock_3', 'Dense_1', 'bias'): (1, 768)
	('DiTBlock_3', 'Dense_1', 'kernel'): (1, 768, 768)
	('DiTBlock_3', 'Dense_2', 'bias'): (1, 768)
	('DiTBlock_3', 'Dense_2', 'kernel'): (1, 768, 768)
	('DiTBlock_3', 'Dense_3', 'bias'): (1, 768)
	('DiTBlock_3', 'Dense_3', 'kernel'): (1, 768, 768)
	('DiTBlock_3', 'Dense_4', 'bias'): (1, 768)
	('DiTBlock_3', 'Dense_4', 'kernel'): (1, 768, 768)
	('DiTBlock_3', 'MlpBlock_0', 'Dense_0', 'bias'): (1, 3072)
	('DiTBlock_3', 'MlpBlock_0', 'Dense_0', 'kernel'): (1, 768, 3072)
	('DiTBlock_3', 'MlpBlock_0', 'Dense_1', 'bias'): (1, 768)
	('DiTBlock_3', 'MlpBlock_0', 'Dense_1', 'kernel'): (1, 3072, 768)
	('DiTBlock_4', 'Dense_0', 'bias'): (1, 4608)
	('DiTBlock_4', 'Dense_0', 'kernel'): (1, 768, 4608)
	('DiTBlock_4', 'Dense_1', 'bias'): (1, 768)
	('DiTBlock_4', 'Dense_1', 'kernel'): (1, 768, 768)
	('DiTBlock_4', 'Dense_2', 'bias'): (1, 768)
	('DiTBlock_4', 'Dense_2', 'kernel'): (1, 768, 768)
	('DiTBlock_4', 'Dense_3', 'bias'): (1, 768)
	('DiTBlock_4', 'Dense_3', 'kernel'): (1, 768, 768)
	('DiTBlock_4', 'Dense_4', 'bias'): (1, 768)
	('DiTBlock_4', 'Dense_4', 'kernel'): (1, 768, 768)
	('DiTBlock_4', 'MlpBlock_0', 'Dense_0', 'bias'): (1, 3072)
	('DiTBlock_4', 'MlpBlock_0', 'Dense_0', 'kernel'): (1, 768, 3072)
	('DiTBlock_4', 'MlpBlock_0', 'Dense_1', 'bias'): (1, 768)
	('DiTBlock_4', 'MlpBlock_0', 'Dense_1', 'kernel'): (1, 3072, 768)
	('DiTBlock_5', 'Dense_0', 'bias'): (1, 4608)
	('DiTBlock_5', 'Dense_0', 'kernel'): (1, 768, 4608)
	('DiTBlock_5', 'Dense_1', 'bias'): (1, 768)
	('DiTBlock_5', 'Dense_1', 'kernel'): (1, 768, 768)
	('DiTBlock_5', 'Dense_2', 'bias'): (1, 768)
	('DiTBlock_5', 'Dense_2', 'kernel'): (1, 768, 768)
	('DiTBlock_5', 'Dense_3', 'bias'): (1, 768)
	('DiTBlock_5', 'Dense_3', 'kernel'): (1, 768, 768)
	('DiTBlock_5', 'Dense_4', 'bias'): (1, 768)
	('DiTBlock_5', 'Dense_4', 'kernel'): (1, 768, 768)
	('DiTBlock_5', 'MlpBlock_0', 'Dense_0', 'bias'): (1, 3072)
	('DiTBlock_5', 'MlpBlock_0', 'Dense_0', 'kernel'): (1, 768, 3072)
	('DiTBlock_5', 'MlpBlock_0', 'Dense_1', 'bias'): (1, 768)
	('DiTBlock_5', 'MlpBlock_0', 'Dense_1', 'kernel'): (1, 3072, 768)
	('DiTBlock_6', 'Dense_0', 'bias'): (1, 4608)
	('DiTBlock_6', 'Dense_0', 'kernel'): (1, 768, 4608)
	('DiTBlock_6', 'Dense_1', 'bias'): (1, 768)
	('DiTBlock_6', 'Dense_1', 'kernel'): (1, 768, 768)
	('DiTBlock_6', 'Dense_2', 'bias'): (1, 768)
	('DiTBlock_6', 'Dense_2', 'kernel'): (1, 768, 768)
	('DiTBlock_6', 'Dense_3', 'bias'): (1, 768)
	('DiTBlock_6', 'Dense_3', 'kernel'): (1, 768, 768)
	('DiTBlock_6', 'Dense_4', 'bias'): (1, 768)
	('DiTBlock_6', 'Dense_4', 'kernel'): (1, 768, 768)
	('DiTBlock_6', 'MlpBlock_0', 'Dense_0', 'bias'): (1, 3072)
	('DiTBlock_6', 'MlpBlock_0', 'Dense_0', 'kernel'): (1, 768, 3072)
	('DiTBlock_6', 'MlpBlock_0', 'Dense_1', 'bias'): (1, 768)
	('DiTBlock_6', 'MlpBlock_0', 'Dense_1', 'kernel'): (1, 3072, 768)
	('DiTBlock_7', 'Dense_0', 'bias'): (1, 4608)
	('DiTBlock_7', 'Dense_0', 'kernel'): (1, 768, 4608)
	('DiTBlock_7', 'Dense_1', 'bias'): (1, 768)
	('DiTBlock_7', 'Dense_1', 'kernel'): (1, 768, 768)
	('DiTBlock_7', 'Dense_2', 'bias'): (1, 768)
	('DiTBlock_7', 'Dense_2', 'kernel'): (1, 768, 768)
	('DiTBlock_7', 'Dense_3', 'bias'): (1, 768)
	('DiTBlock_7', 'Dense_3', 'kernel'): (1, 768, 768)
	('DiTBlock_7', 'Dense_4', 'bias'): (1, 768)
	('DiTBlock_7', 'Dense_4', 'kernel'): (1, 768, 768)
	('DiTBlock_7', 'MlpBlock_0', 'Dense_0', 'bias'): (1, 3072)
	('DiTBlock_7', 'MlpBlock_0', 'Dense_0', 'kernel'): (1, 768, 3072)
	('DiTBlock_7', 'MlpBlock_0', 'Dense_1', 'bias'): (1, 768)
	('DiTBlock_7', 'MlpBlock_0', 'Dense_1', 'kernel'): (1, 3072, 768)
	('DiTBlock_8', 'Dense_0', 'bias'): (1, 4608)
	('DiTBlock_8', 'Dense_0', 'kernel'): (1, 768, 4608)
	('DiTBlock_8', 'Dense_1', 'bias'): (1, 768)
	('DiTBlock_8', 'Dense_1', 'kernel'): (1, 768, 768)
	('DiTBlock_8', 'Dense_2', 'bias'): (1, 768)
	('DiTBlock_8', 'Dense_2', 'kernel'): (1, 768, 768)
	('DiTBlock_8', 'Dense_3', 'bias'): (1, 768)
	('DiTBlock_8', 'Dense_3', 'kernel'): (1, 768, 768)
	('DiTBlock_8', 'Dense_4', 'bias'): (1, 768)
	('DiTBlock_8', 'Dense_4', 'kernel'): (1, 768, 768)
	('DiTBlock_8', 'MlpBlock_0', 'Dense_0', 'bias'): (1, 3072)
	('DiTBlock_8', 'MlpBlock_0', 'Dense_0', 'kernel'): (1, 768, 3072)
	('DiTBlock_8', 'MlpBlock_0', 'Dense_1', 'bias'): (1, 768)
	('DiTBlock_8', 'MlpBlock_0', 'Dense_1', 'kernel'): (1, 3072, 768)
	('DiTBlock_9', 'Dense_0', 'bias'): (1, 4608)
	('DiTBlock_9', 'Dense_0', 'kernel'): (1, 768, 4608)
	('DiTBlock_9', 'Dense_1', 'bias'): (1, 768)
	('DiTBlock_9', 'Dense_1', 'kernel'): (1, 768, 768)
	('DiTBlock_9', 'Dense_2', 'bias'): (1, 768)
	('DiTBlock_9', 'Dense_2', 'kernel'): (1, 768, 768)
	('DiTBlock_9', 'Dense_3', 'bias'): (1, 768)
	('DiTBlock_9', 'Dense_3', 'kernel'): (1, 768, 768)
	('DiTBlock_9', 'Dense_4', 'bias'): (1, 768)
	('DiTBlock_9', 'Dense_4', 'kernel'): (1, 768, 768)
	('DiTBlock_9', 'MlpBlock_0', 'Dense_0', 'bias'): (1, 3072)
	('DiTBlock_9', 'MlpBlock_0', 'Dense_0', 'kernel'): (1, 768, 3072)
	('DiTBlock_9', 'MlpBlock_0', 'Dense_1', 'bias'): (1, 768)
	('DiTBlock_9', 'MlpBlock_0', 'Dense_1', 'kernel'): (1, 3072, 768)
	('Embed_0', 'embedding'): (1, 256, 1)
	('FinalLayer_0', 'Dense_0', 'bias'): (1, 1536)
	('FinalLayer_0', 'Dense_0', 'kernel'): (1, 768, 1536)
	('FinalLayer_0', 'Dense_1', 'bias'): (1, 16)
	('FinalLayer_0', 'Dense_1', 'kernel'): (1, 768, 16)
	('LabelEmbedder_0', 'Embed_0', 'embedding'): (1, 1001, 768)
	('PatchEmbed_0', 'Conv_0', 'bias'): (1, 768)
	('PatchEmbed_0', 'Conv_0', 'kernel'): (1, 2, 2, 4, 768)
	('TimestepEmbedder_0', 'Dense_0', 'bias'): (1, 768)
	('TimestepEmbedder_0', 'Dense_0', 'kernel'): (1, 256, 768)
	('TimestepEmbedder_0', 'Dense_1', 'bias'): (1, 768)
	('TimestepEmbedder_0', 'Dense_1', 'kernel'): (1, 768, 768)
	('TimestepEmbedder_1', 'Dense_0', 'bias'): (1, 768)
	('TimestepEmbedder_1', 'Dense_0', 'kernel'): (1, 256, 768)
	('TimestepEmbedder_1', 'Dense_1', 'bias'): (1, 768)
	('TimestepEmbedder_1', 'Dense_1', 'kernel'): (1, 768, 768)

	parameter shapes:
	('DiTBlock_0', 'Dense_0', 'bias'): (1, 4608)
	('DiTBlock_0', 'Dense_0', 'kernel'): (1, 768, 4608)
	('DiTBlock_0', 'Dense_1', 'bias'): (1, 768)
	('DiTBlock_0', 'Dense_1', 'kernel'): (1, 768, 768)
	('DiTBlock_0', 'Dense_2', 'bias'): (1, 768)
	('DiTBlock_0', 'Dense_2', 'kernel'): (1, 768, 768)
	('DiTBlock_0', 'Dense_3', 'bias'): (1, 768)
	('DiTBlock_0', 'Dense_3', 'kernel'): (1, 768, 768)
	('DiTBlock_0', 'Dense_4', 'bias'): (1, 768)
	('DiTBlock_0', 'Dense_4', 'kernel'): (1, 768, 768)
	('DiTBlock_0', 'MlpBlock_0', 'Dense_0', 'bias'): (1, 3072)
	('DiTBlock_0', 'MlpBlock_0', 'Dense_0', 'kernel'): (1, 768, 3072)
	('DiTBlock_0', 'MlpBlock_0', 'Dense_1', 'bias'): (1, 768)
	('DiTBlock_0', 'MlpBlock_0', 'Dense_1', 'kernel'): (1, 3072, 768)
	('DiTBlock_1', 'Dense_0', 'bias'): (1, 4608)
	('DiTBlock_1', 'Dense_0', 'kernel'): (1, 768, 4608)
	('DiTBlock_1', 'Dense_1', 'bias'): (1, 768)
	('DiTBlock_1', 'Dense_1', 'kernel'): (1, 768, 768)
	('DiTBlock_1', 'Dense_2', 'bias'): (1, 768)
	('DiTBlock_1', 'Dense_2', 'kernel'): (1, 768, 768)
	('DiTBlock_1', 'Dense_3', 'bias'): (1, 768)
	('DiTBlock_1', 'Dense_3', 'kernel'): (1, 768, 768)
	('DiTBlock_1', 'Dense_4', 'bias'): (1, 768)
	('DiTBlock_1', 'Dense_4', 'kernel'): (1, 768, 768)
	('DiTBlock_1', 'MlpBlock_0', 'Dense_0', 'bias'): (1, 3072)
	('DiTBlock_1', 'MlpBlock_0', 'Dense_0', 'kernel'): (1, 768, 3072)
	('DiTBlock_1', 'MlpBlock_0', 'Dense_1', 'bias'): (1, 768)
	('DiTBlock_1', 'MlpBlock_0', 'Dense_1', 'kernel'): (1, 3072, 768)
	('DiTBlock_10', 'Dense_0', 'bias'): (1, 4608)
	('DiTBlock_10', 'Dense_0', 'kernel'): (1, 768, 4608)
	('DiTBlock_10', 'Dense_1', 'bias'): (1, 768)
	('DiTBlock_10', 'Dense_1', 'kernel'): (1, 768, 768)
	('DiTBlock_10', 'Dense_2', 'bias'): (1, 768)
	('DiTBlock_10', 'Dense_2', 'kernel'): (1, 768, 768)
	('DiTBlock_10', 'Dense_3', 'bias'): (1, 768)
	('DiTBlock_10', 'Dense_3', 'kernel'): (1, 768, 768)
	('DiTBlock_10', 'Dense_4', 'bias'): (1, 768)
	('DiTBlock_10', 'Dense_4', 'kernel'): (1, 768, 768)
	('DiTBlock_10', 'MlpBlock_0', 'Dense_0', 'bias'): (1, 3072)
	('DiTBlock_10', 'MlpBlock_0', 'Dense_0', 'kernel'): (1, 768, 3072)
	('DiTBlock_10', 'MlpBlock_0', 'Dense_1', 'bias'): (1, 768)
	('DiTBlock_10', 'MlpBlock_0', 'Dense_1', 'kernel'): (1, 3072, 768)
	('DiTBlock_11', 'Dense_0', 'bias'): (1, 4608)
	('DiTBlock_11', 'Dense_0', 'kernel'): (1, 768, 4608)
	('DiTBlock_11', 'Dense_1', 'bias'): (1, 768)
	('DiTBlock_11', 'Dense_1', 'kernel'): (1, 768, 768)
	('DiTBlock_11', 'Dense_2', 'bias'): (1, 768)
	('DiTBlock_11', 'Dense_2', 'kernel'): (1, 768, 768)
	('DiTBlock_11', 'Dense_3', 'bias'): (1, 768)
	('DiTBlock_11', 'Dense_3', 'kernel'): (1, 768, 768)
	('DiTBlock_11', 'Dense_4', 'bias'): (1, 768)
	('DiTBlock_11', 'Dense_4', 'kernel'): (1, 768, 768)
	('DiTBlock_11', 'MlpBlock_0', 'Dense_0', 'bias'): (1, 3072)
	('DiTBlock_11', 'MlpBlock_0', 'Dense_0', 'kernel'): (1, 768, 3072)
	('DiTBlock_11', 'MlpBlock_0', 'Dense_1', 'bias'): (1, 768)
	('DiTBlock_11', 'MlpBlock_0', 'Dense_1', 'kernel'): (1, 3072, 768)
	('DiTBlock_2', 'Dense_0', 'bias'): (1, 4608)
	('DiTBlock_2', 'Dense_0', 'kernel'): (1, 768, 4608)
	('DiTBlock_2', 'Dense_1', 'bias'): (1, 768)
	('DiTBlock_2', 'Dense_1', 'kernel'): (1, 768, 768)
	('DiTBlock_2', 'Dense_2', 'bias'): (1, 768)
	('DiTBlock_2', 'Dense_2', 'kernel'): (1, 768, 768)
	('DiTBlock_2', 'Dense_3', 'bias'): (1, 768)
	('DiTBlock_2', 'Dense_3', 'kernel'): (1, 768, 768)
	('DiTBlock_2', 'Dense_4', 'bias'): (1, 768)
	('DiTBlock_2', 'Dense_4', 'kernel'): (1, 768, 768)
	('DiTBlock_2', 'MlpBlock_0', 'Dense_0', 'bias'): (1, 3072)
	('DiTBlock_2', 'MlpBlock_0', 'Dense_0', 'kernel'): (1, 768, 3072)
	('DiTBlock_2', 'MlpBlock_0', 'Dense_1', 'bias'): (1, 768)
	('DiTBlock_2', 'MlpBlock_0', 'Dense_1', 'kernel'): (1, 3072, 768)
	('DiTBlock_3', 'Dense_0', 'bias'): (1, 4608)
	('DiTBlock_3', 'Dense_0', 'kernel'): (1, 768, 4608)
	('DiTBlock_3', 'Dense_1', 'bias'): (1, 768)
	('DiTBlock_3', 'Dense_1', 'kernel'): (1, 768, 768)
	('DiTBlock_3', 'Dense_2', 'bias'): (1, 768)
	('DiTBlock_3', 'Dense_2', 'kernel'): (1, 768, 768)
	('DiTBlock_3', 'Dense_3', 'bias'): (1, 768)
	('DiTBlock_3', 'Dense_3', 'kernel'): (1, 768, 768)
	('DiTBlock_3', 'Dense_4', 'bias'): (1, 768)
	('DiTBlock_3', 'Dense_4', 'kernel'): (1, 768, 768)
	('DiTBlock_3', 'MlpBlock_0', 'Dense_0', 'bias'): (1, 3072)
	('DiTBlock_3', 'MlpBlock_0', 'Dense_0', 'kernel'): (1, 768, 3072)
	('DiTBlock_3', 'MlpBlock_0', 'Dense_1', 'bias'): (1, 768)
	('DiTBlock_3', 'MlpBlock_0', 'Dense_1', 'kernel'): (1, 3072, 768)
	('DiTBlock_4', 'Dense_0', 'bias'): (1, 4608)
	('DiTBlock_4', 'Dense_0', 'kernel'): (1, 768, 4608)
	('DiTBlock_4', 'Dense_1', 'bias'): (1, 768)
	('DiTBlock_4', 'Dense_1', 'kernel'): (1, 768, 768)
	('DiTBlock_4', 'Dense_2', 'bias'): (1, 768)
	('DiTBlock_4', 'Dense_2', 'kernel'): (1, 768, 768)
	('DiTBlock_4', 'Dense_3', 'bias'): (1, 768)
	('DiTBlock_4', 'Dense_3', 'kernel'): (1, 768, 768)
	('DiTBlock_4', 'Dense_4', 'bias'): (1, 768)
	('DiTBlock_4', 'Dense_4', 'kernel'): (1, 768, 768)
	('DiTBlock_4', 'MlpBlock_0', 'Dense_0', 'bias'): (1, 3072)
	('DiTBlock_4', 'MlpBlock_0', 'Dense_0', 'kernel'): (1, 768, 3072)
	('DiTBlock_4', 'MlpBlock_0', 'Dense_1', 'bias'): (1, 768)
	('DiTBlock_4', 'MlpBlock_0', 'Dense_1', 'kernel'): (1, 3072, 768)
	('DiTBlock_5', 'Dense_0', 'bias'): (1, 4608)
	('DiTBlock_5', 'Dense_0', 'kernel'): (1, 768, 4608)
	('DiTBlock_5', 'Dense_1', 'bias'): (1, 768)
	('DiTBlock_5', 'Dense_1', 'kernel'): (1, 768, 768)
	('DiTBlock_5', 'Dense_2', 'bias'): (1, 768)
	('DiTBlock_5', 'Dense_2', 'kernel'): (1, 768, 768)
	('DiTBlock_5', 'Dense_3', 'bias'): (1, 768)
	('DiTBlock_5', 'Dense_3', 'kernel'): (1, 768, 768)
	('DiTBlock_5', 'Dense_4', 'bias'): (1, 768)
	('DiTBlock_5', 'Dense_4', 'kernel'): (1, 768, 768)
	('DiTBlock_5', 'MlpBlock_0', 'Dense_0', 'bias'): (1, 3072)
	('DiTBlock_5', 'MlpBlock_0', 'Dense_0', 'kernel'): (1, 768, 3072)
	('DiTBlock_5', 'MlpBlock_0', 'Dense_1', 'bias'): (1, 768)
	('DiTBlock_5', 'MlpBlock_0', 'Dense_1', 'kernel'): (1, 3072, 768)
	('DiTBlock_6', 'Dense_0', 'bias'): (1, 4608)
	('DiTBlock_6', 'Dense_0', 'kernel'): (1, 768, 4608)
	('DiTBlock_6', 'Dense_1', 'bias'): (1, 768)
	('DiTBlock_6', 'Dense_1', 'kernel'): (1, 768, 768)
	('DiTBlock_6', 'Dense_2', 'bias'): (1, 768)
	('DiTBlock_6', 'Dense_2', 'kernel'): (1, 768, 768)
	('DiTBlock_6', 'Dense_3', 'bias'): (1, 768)
	('DiTBlock_6', 'Dense_3', 'kernel'): (1, 768, 768)
	('DiTBlock_6', 'Dense_4', 'bias'): (1, 768)
	('DiTBlock_6', 'Dense_4', 'kernel'): (1, 768, 768)
	('DiTBlock_6', 'MlpBlock_0', 'Dense_0', 'bias'): (1, 3072)
	('DiTBlock_6', 'MlpBlock_0', 'Dense_0', 'kernel'): (1, 768, 3072)
	('DiTBlock_6', 'MlpBlock_0', 'Dense_1', 'bias'): (1, 768)
	('DiTBlock_6', 'MlpBlock_0', 'Dense_1', 'kernel'): (1, 3072, 768)
	('DiTBlock_7', 'Dense_0', 'bias'): (1, 4608)
	('DiTBlock_7', 'Dense_0', 'kernel'): (1, 768, 4608)
	('DiTBlock_7', 'Dense_1', 'bias'): (1, 768)
	('DiTBlock_7', 'Dense_1', 'kernel'): (1, 768, 768)
	('DiTBlock_7', 'Dense_2', 'bias'): (1, 768)
	('DiTBlock_7', 'Dense_2', 'kernel'): (1, 768, 768)
	('DiTBlock_7', 'Dense_3', 'bias'): (1, 768)
	('DiTBlock_7', 'Dense_3', 'kernel'): (1, 768, 768)
	('DiTBlock_7', 'Dense_4', 'bias'): (1, 768)
	('DiTBlock_7', 'Dense_4', 'kernel'): (1, 768, 768)
	('DiTBlock_7', 'MlpBlock_0', 'Dense_0', 'bias'): (1, 3072)
	('DiTBlock_7', 'MlpBlock_0', 'Dense_0', 'kernel'): (1, 768, 3072)
	('DiTBlock_7', 'MlpBlock_0', 'Dense_1', 'bias'): (1, 768)
	('DiTBlock_7', 'MlpBlock_0', 'Dense_1', 'kernel'): (1, 3072, 768)
	('DiTBlock_8', 'Dense_0', 'bias'): (1, 4608)
	('DiTBlock_8', 'Dense_0', 'kernel'): (1, 768, 4608)
	('DiTBlock_8', 'Dense_1', 'bias'): (1, 768)
	('DiTBlock_8', 'Dense_1', 'kernel'): (1, 768, 768)
	('DiTBlock_8', 'Dense_2', 'bias'): (1, 768)
	('DiTBlock_8', 'Dense_2', 'kernel'): (1, 768, 768)
	('DiTBlock_8', 'Dense_3', 'bias'): (1, 768)
	('DiTBlock_8', 'Dense_3', 'kernel'): (1, 768, 768)
	('DiTBlock_8', 'Dense_4', 'bias'): (1, 768)
	('DiTBlock_8', 'Dense_4', 'kernel'): (1, 768, 768)
	('DiTBlock_8', 'MlpBlock_0', 'Dense_0', 'bias'): (1, 3072)
	('DiTBlock_8', 'MlpBlock_0', 'Dense_0', 'kernel'): (1, 768, 3072)
	('DiTBlock_8', 'MlpBlock_0', 'Dense_1', 'bias'): (1, 768)
	('DiTBlock_8', 'MlpBlock_0', 'Dense_1', 'kernel'): (1, 3072, 768)
	('DiTBlock_9', 'Dense_0', 'bias'): (1, 4608)
	('DiTBlock_9', 'Dense_0', 'kernel'): (1, 768, 4608)
	('DiTBlock_9', 'Dense_1', 'bias'): (1, 768)
	('DiTBlock_9', 'Dense_1', 'kernel'): (1, 768, 768)
	('DiTBlock_9', 'Dense_2', 'bias'): (1, 768)
	('DiTBlock_9', 'Dense_2', 'kernel'): (1, 768, 768)
	('DiTBlock_9', 'Dense_3', 'bias'): (1, 768)
	('DiTBlock_9', 'Dense_3', 'kernel'): (1, 768, 768)
	('DiTBlock_9', 'Dense_4', 'bias'): (1, 768)
	('DiTBlock_9', 'Dense_4', 'kernel'): (1, 768, 768)
	('DiTBlock_9', 'MlpBlock_0', 'Dense_0', 'bias'): (1, 3072)
	('DiTBlock_9', 'MlpBlock_0', 'Dense_0', 'kernel'): (1, 768, 3072)
	('DiTBlock_9', 'MlpBlock_0', 'Dense_1', 'bias'): (1, 768)
	('DiTBlock_9', 'MlpBlock_0', 'Dense_1', 'kernel'): (1, 3072, 768)
	('Embed_0', 'embedding'): (1, 256, 1)
	('FinalLayer_0', 'Dense_0', 'bias'): (1, 1536)
	('FinalLayer_0', 'Dense_0', 'kernel'): (1, 768, 1536)
	('FinalLayer_0', 'Dense_1', 'bias'): (1, 16)
	('FinalLayer_0', 'Dense_1', 'kernel'): (1, 768, 16)
	('LabelEmbedder_0', 'Embed_0', 'embedding'): (1, 1001, 768)
	('PatchEmbed_0', 'Conv_0', 'bias'): (1, 768)
	('PatchEmbed_0', 'Conv_0', 'kernel'): (1, 2, 2, 4, 768)
	('TimestepEmbedder_0', 'Dense_0', 'bias'): (1, 768)
	('TimestepEmbedder_0', 'Dense_0', 'kernel'): (1, 256, 768)
	('TimestepEmbedder_0', 'Dense_1', 'bias'): (1, 768)
	('TimestepEmbedder_0', 'Dense_1', 'kernel'): (1, 768, 768)
	('TimestepEmbedder_1', 'Dense_0', 'bias'): (1, 768)
	('TimestepEmbedder_1', 'Dense_0', 'kernel'): (1, 256, 768)
	('TimestepEmbedder_1', 'Dense_1', 'bias'): (1, 768)
	('TimestepEmbedder_1', 'Dense_1', 'kernel'): (1, 768, 768)

	parameter shapes:
	('DiTBlock_0', 'Dense_0', 'bias'): (1, 4608)
	('DiTBlock_0', 'Dense_0', 'kernel'): (1, 768, 4608)
	('DiTBlock_0', 'Dense_1', 'bias'): (1, 768)
	('DiTBlock_0', 'Dense_1', 'kernel'): (1, 768, 768)
	('DiTBlock_0', 'Dense_2', 'bias'): (1, 768)
	('DiTBlock_0', 'Dense_2', 'kernel'): (1, 768, 768)
	('DiTBlock_0', 'Dense_3', 'bias'): (1, 768)
	('DiTBlock_0', 'Dense_3', 'kernel'): (1, 768, 768)
	('DiTBlock_0', 'Dense_4', 'bias'): (1, 768)
	('DiTBlock_0', 'Dense_4', 'kernel'): (1, 768, 768)
	('DiTBlock_0', 'MlpBlock_0', 'Dense_0', 'bias'): (1, 3072)
	('DiTBlock_0', 'MlpBlock_0', 'Dense_0', 'kernel'): (1, 768, 3072)
	('DiTBlock_0', 'MlpBlock_0', 'Dense_1', 'bias'): (1, 768)
	('DiTBlock_0', 'MlpBlock_0', 'Dense_1', 'kernel'): (1, 3072, 768)
	('DiTBlock_1', 'Dense_0', 'bias'): (1, 4608)
	('DiTBlock_1', 'Dense_0', 'kernel'): (1, 768, 4608)
	('DiTBlock_1', 'Dense_1', 'bias'): (1, 768)
	('DiTBlock_1', 'Dense_1', 'kernel'): (1, 768, 768)
	('DiTBlock_1', 'Dense_2', 'bias'): (1, 768)
	('DiTBlock_1', 'Dense_2', 'kernel'): (1, 768, 768)
	('DiTBlock_1', 'Dense_3', 'bias'): (1, 768)
	('DiTBlock_1', 'Dense_3', 'kernel'): (1, 768, 768)
	('DiTBlock_1', 'Dense_4', 'bias'): (1, 768)
	('DiTBlock_1', 'Dense_4', 'kernel'): (1, 768, 768)
	('DiTBlock_1', 'MlpBlock_0', 'Dense_0', 'bias'): (1, 3072)
	('DiTBlock_1', 'MlpBlock_0', 'Dense_0', 'kernel'): (1, 768, 3072)
	('DiTBlock_1', 'MlpBlock_0', 'Dense_1', 'bias'): (1, 768)
	('DiTBlock_1', 'MlpBlock_0', 'Dense_1', 'kernel'): (1, 3072, 768)
	('DiTBlock_10', 'Dense_0', 'bias'): (1, 4608)
	('DiTBlock_10', 'Dense_0', 'kernel'): (1, 768, 4608)
	('DiTBlock_10', 'Dense_1', 'bias'): (1, 768)
	('DiTBlock_10', 'Dense_1', 'kernel'): (1, 768, 768)
	('DiTBlock_10', 'Dense_2', 'bias'): (1, 768)
	('DiTBlock_10', 'Dense_2', 'kernel'): (1, 768, 768)
	('DiTBlock_10', 'Dense_3', 'bias'): (1, 768)
	('DiTBlock_10', 'Dense_3', 'kernel'): (1, 768, 768)
	('DiTBlock_10', 'Dense_4', 'bias'): (1, 768)
	('DiTBlock_10', 'Dense_4', 'kernel'): (1, 768, 768)
	('DiTBlock_10', 'MlpBlock_0', 'Dense_0', 'bias'): (1, 3072)
	('DiTBlock_10', 'MlpBlock_0', 'Dense_0', 'kernel'): (1, 768, 3072)
	('DiTBlock_10', 'MlpBlock_0', 'Dense_1', 'bias'): (1, 768)
	('DiTBlock_10', 'MlpBlock_0', 'Dense_1', 'kernel'): (1, 3072, 768)
	('DiTBlock_11', 'Dense_0', 'bias'): (1, 4608)
	('DiTBlock_11', 'Dense_0', 'kernel'): (1, 768, 4608)
	('DiTBlock_11', 'Dense_1', 'bias'): (1, 768)
	('DiTBlock_11', 'Dense_1', 'kernel'): (1, 768, 768)
	('DiTBlock_11', 'Dense_2', 'bias'): (1, 768)
	('DiTBlock_11', 'Dense_2', 'kernel'): (1, 768, 768)
	('DiTBlock_11', 'Dense_3', 'bias'): (1, 768)
	('DiTBlock_11', 'Dense_3', 'kernel'): (1, 768, 768)
	('DiTBlock_11', 'Dense_4', 'bias'): (1, 768)
	('DiTBlock_11', 'Dense_4', 'kernel'): (1, 768, 768)
	('DiTBlock_11', 'MlpBlock_0', 'Dense_0', 'bias'): (1, 3072)
	('DiTBlock_11', 'MlpBlock_0', 'Dense_0', 'kernel'): (1, 768, 3072)
	('DiTBlock_11', 'MlpBlock_0', 'Dense_1', 'bias'): (1, 768)
	('DiTBlock_11', 'MlpBlock_0', 'Dense_1', 'kernel'): (1, 3072, 768)
	('DiTBlock_2', 'Dense_0', 'bias'): (1, 4608)
	('DiTBlock_2', 'Dense_0', 'kernel'): (1, 768, 4608)
	('DiTBlock_2', 'Dense_1', 'bias'): (1, 768)
	('DiTBlock_2', 'Dense_1', 'kernel'): (1, 768, 768)
	('DiTBlock_2', 'Dense_2', 'bias'): (1, 768)
	('DiTBlock_2', 'Dense_2', 'kernel'): (1, 768, 768)
	('DiTBlock_2', 'Dense_3', 'bias'): (1, 768)
	('DiTBlock_2', 'Dense_3', 'kernel'): (1, 768, 768)
	('DiTBlock_2', 'Dense_4', 'bias'): (1, 768)
	('DiTBlock_2', 'Dense_4', 'kernel'): (1, 768, 768)
	('DiTBlock_2', 'MlpBlock_0', 'Dense_0', 'bias'): (1, 3072)
	('DiTBlock_2', 'MlpBlock_0', 'Dense_0', 'kernel'): (1, 768, 3072)
	('DiTBlock_2', 'MlpBlock_0', 'Dense_1', 'bias'): (1, 768)
	('DiTBlock_2', 'MlpBlock_0', 'Dense_1', 'kernel'): (1, 3072, 768)
	('DiTBlock_3', 'Dense_0', 'bias'): (1, 4608)
	('DiTBlock_3', 'Dense_0', 'kernel'): (1, 768, 4608)
	('DiTBlock_3', 'Dense_1', 'bias'): (1, 768)
	('DiTBlock_3', 'Dense_1', 'kernel'): (1, 768, 768)
	('DiTBlock_3', 'Dense_2', 'bias'): (1, 768)
	('DiTBlock_3', 'Dense_2', 'kernel'): (1, 768, 768)
	('DiTBlock_3', 'Dense_3', 'bias'): (1, 768)
	('DiTBlock_3', 'Dense_3', 'kernel'): (1, 768, 768)
	('DiTBlock_3', 'Dense_4', 'bias'): (1, 768)
	('DiTBlock_3', 'Dense_4', 'kernel'): (1, 768, 768)
	('DiTBlock_3', 'MlpBlock_0', 'Dense_0', 'bias'): (1, 3072)
	('DiTBlock_3', 'MlpBlock_0', 'Dense_0', 'kernel'): (1, 768, 3072)
	('DiTBlock_3', 'MlpBlock_0', 'Dense_1', 'bias'): (1, 768)
	('DiTBlock_3', 'MlpBlock_0', 'Dense_1', 'kernel'): (1, 3072, 768)
	('DiTBlock_4', 'Dense_0', 'bias'): (1, 4608)
	('DiTBlock_4', 'Dense_0', 'kernel'): (1, 768, 4608)
	('DiTBlock_4', 'Dense_1', 'bias'): (1, 768)
	('DiTBlock_4', 'Dense_1', 'kernel'): (1, 768, 768)
	('DiTBlock_4', 'Dense_2', 'bias'): (1, 768)
	('DiTBlock_4', 'Dense_2', 'kernel'): (1, 768, 768)
	('DiTBlock_4', 'Dense_3', 'bias'): (1, 768)
	('DiTBlock_4', 'Dense_3', 'kernel'): (1, 768, 768)
	('DiTBlock_4', 'Dense_4', 'bias'): (1, 768)
	('DiTBlock_4', 'Dense_4', 'kernel'): (1, 768, 768)
	('DiTBlock_4', 'MlpBlock_0', 'Dense_0', 'bias'): (1, 3072)
	('DiTBlock_4', 'MlpBlock_0', 'Dense_0', 'kernel'): (1, 768, 3072)
	('DiTBlock_4', 'MlpBlock_0', 'Dense_1', 'bias'): (1, 768)
	('DiTBlock_4', 'MlpBlock_0', 'Dense_1', 'kernel'): (1, 3072, 768)
	('DiTBlock_5', 'Dense_0', 'bias'): (1, 4608)
	('DiTBlock_5', 'Dense_0', 'kernel'): (1, 768, 4608)
	('DiTBlock_5', 'Dense_1', 'bias'): (1, 768)
	('DiTBlock_5', 'Dense_1', 'kernel'): (1, 768, 768)
	('DiTBlock_5', 'Dense_2', 'bias'): (1, 768)
	('DiTBlock_5', 'Dense_2', 'kernel'): (1, 768, 768)
	('DiTBlock_5', 'Dense_3', 'bias'): (1, 768)
	('DiTBlock_5', 'Dense_3', 'kernel'): (1, 768, 768)
	('DiTBlock_5', 'Dense_4', 'bias'): (1, 768)
	('DiTBlock_5', 'Dense_4', 'kernel'): (1, 768, 768)
	('DiTBlock_5', 'MlpBlock_0', 'Dense_0', 'bias'): (1, 3072)
	('DiTBlock_5', 'MlpBlock_0', 'Dense_0', 'kernel'): (1, 768, 3072)
	('DiTBlock_5', 'MlpBlock_0', 'Dense_1', 'bias'): (1, 768)
	('DiTBlock_5', 'MlpBlock_0', 'Dense_1', 'kernel'): (1, 3072, 768)
	('DiTBlock_6', 'Dense_0', 'bias'): (1, 4608)
	('DiTBlock_6', 'Dense_0', 'kernel'): (1, 768, 4608)
	('DiTBlock_6', 'Dense_1', 'bias'): (1, 768)
	('DiTBlock_6', 'Dense_1', 'kernel'): (1, 768, 768)
	('DiTBlock_6', 'Dense_2', 'bias'): (1, 768)
	('DiTBlock_6', 'Dense_2', 'kernel'): (1, 768, 768)
	('DiTBlock_6', 'Dense_3', 'bias'): (1, 768)
	('DiTBlock_6', 'Dense_3', 'kernel'): (1, 768, 768)
	('DiTBlock_6', 'Dense_4', 'bias'): (1, 768)
	('DiTBlock_6', 'Dense_4', 'kernel'): (1, 768, 768)
	('DiTBlock_6', 'MlpBlock_0', 'Dense_0', 'bias'): (1, 3072)
	('DiTBlock_6', 'MlpBlock_0', 'Dense_0', 'kernel'): (1, 768, 3072)
	('DiTBlock_6', 'MlpBlock_0', 'Dense_1', 'bias'): (1, 768)
	('DiTBlock_6', 'MlpBlock_0', 'Dense_1', 'kernel'): (1, 3072, 768)
	('DiTBlock_7', 'Dense_0', 'bias'): (1, 4608)
	('DiTBlock_7', 'Dense_0', 'kernel'): (1, 768, 4608)
	('DiTBlock_7', 'Dense_1', 'bias'): (1, 768)
	('DiTBlock_7', 'Dense_1', 'kernel'): (1, 768, 768)
	('DiTBlock_7', 'Dense_2', 'bias'): (1, 768)
	('DiTBlock_7', 'Dense_2', 'kernel'): (1, 768, 768)
	('DiTBlock_7', 'Dense_3', 'bias'): (1, 768)
	('DiTBlock_7', 'Dense_3', 'kernel'): (1, 768, 768)
	('DiTBlock_7', 'Dense_4', 'bias'): (1, 768)
	('DiTBlock_7', 'Dense_4', 'kernel'): (1, 768, 768)
	('DiTBlock_7', 'MlpBlock_0', 'Dense_0', 'bias'): (1, 3072)
	('DiTBlock_7', 'MlpBlock_0', 'Dense_0', 'kernel'): (1, 768, 3072)
	('DiTBlock_7', 'MlpBlock_0', 'Dense_1', 'bias'): (1, 768)
	('DiTBlock_7', 'MlpBlock_0', 'Dense_1', 'kernel'): (1, 3072, 768)
	('DiTBlock_8', 'Dense_0', 'bias'): (1, 4608)
	('DiTBlock_8', 'Dense_0', 'kernel'): (1, 768, 4608)
	('DiTBlock_8', 'Dense_1', 'bias'): (1, 768)
	('DiTBlock_8', 'Dense_1', 'kernel'): (1, 768, 768)
	('DiTBlock_8', 'Dense_2', 'bias'): (1, 768)
	('DiTBlock_8', 'Dense_2', 'kernel'): (1, 768, 768)
	('DiTBlock_8', 'Dense_3', 'bias'): (1, 768)
	('DiTBlock_8', 'Dense_3', 'kernel'): (1, 768, 768)
	('DiTBlock_8', 'Dense_4', 'bias'): (1, 768)
	('DiTBlock_8', 'Dense_4', 'kernel'): (1, 768, 768)
	('DiTBlock_8', 'MlpBlock_0', 'Dense_0', 'bias'): (1, 3072)
	('DiTBlock_8', 'MlpBlock_0', 'Dense_0', 'kernel'): (1, 768, 3072)
	('DiTBlock_8', 'MlpBlock_0', 'Dense_1', 'bias'): (1, 768)
	('DiTBlock_8', 'MlpBlock_0', 'Dense_1', 'kernel'): (1, 3072, 768)
	('DiTBlock_9', 'Dense_0', 'bias'): (1, 4608)
	('DiTBlock_9', 'Dense_0', 'kernel'): (1, 768, 4608)
	('DiTBlock_9', 'Dense_1', 'bias'): (1, 768)
	('DiTBlock_9', 'Dense_1', 'kernel'): (1, 768, 768)
	('DiTBlock_9', 'Dense_2', 'bias'): (1, 768)
	('DiTBlock_9', 'Dense_2', 'kernel'): (1, 768, 768)
	('DiTBlock_9', 'Dense_3', 'bias'): (1, 768)
	('DiTBlock_9', 'Dense_3', 'kernel'): (1, 768, 768)
	('DiTBlock_9', 'Dense_4', 'bias'): (1, 768)
	('DiTBlock_9', 'Dense_4', 'kernel'): (1, 768, 768)
	('DiTBlock_9', 'MlpBlock_0', 'Dense_0', 'bias'): (1, 3072)
	('DiTBlock_9', 'MlpBlock_0', 'Dense_0', 'kernel'): (1, 768, 3072)
	('DiTBlock_9', 'MlpBlock_0', 'Dense_1', 'bias'): (1, 768)
	('DiTBlock_9', 'MlpBlock_0', 'Dense_1', 'kernel'): (1, 3072, 768)
	('Embed_0', 'embedding'): (1, 256, 1)
	('FinalLayer_0', 'Dense_0', 'bias'): (1, 1536)
	('FinalLayer_0', 'Dense_0', 'kernel'): (1, 768, 1536)
	('FinalLayer_0', 'Dense_1', 'bias'): (1, 16)
	('FinalLayer_0', 'Dense_1', 'kernel'): (1, 768, 16)
	('LabelEmbedder_0', 'Embed_0', 'embedding'): (1, 1001, 768)
	('PatchEmbed_0', 'Conv_0', 'bias'): (1, 768)
	('PatchEmbed_0', 'Conv_0', 'kernel'): (1, 2, 2, 4, 768)
	('TimestepEmbedder_0', 'Dense_0', 'bias'): (1, 768)
	('TimestepEmbedder_0', 'Dense_0', 'kernel'): (1, 256, 768)
	('TimestepEmbedder_0', 'Dense_1', 'bias'): (1, 768)
	('TimestepEmbedder_0', 'Dense_1', 'kernel'): (1, 768, 768)
	('TimestepEmbedder_1', 'Dense_0', 'bias'): (1, 768)
	('TimestepEmbedder_1', 'Dense_0', 'kernel'): (1, 256, 768)
	('TimestepEmbedder_1', 'Dense_1', 'bias'): (1, 768)
	('TimestepEmbedder_1', 'Dense_1', 'kernel'): (1, 768, 768)

	parameter shapes:
	('DiTBlock_0', 'Dense_0', 'bias'): (1, 4608)
	('DiTBlock_0', 'Dense_0', 'kernel'): (1, 768, 4608)
	('DiTBlock_0', 'Dense_1', 'bias'): (1, 768)
	('DiTBlock_0', 'Dense_1', 'kernel'): (1, 768, 768)
	('DiTBlock_0', 'Dense_2', 'bias'): (1, 768)
	('DiTBlock_0', 'Dense_2', 'kernel'): (1, 768, 768)
	('DiTBlock_0', 'Dense_3', 'bias'): (1, 768)
	('DiTBlock_0', 'Dense_3', 'kernel'): (1, 768, 768)
	('DiTBlock_0', 'Dense_4', 'bias'): (1, 768)
	('DiTBlock_0', 'Dense_4', 'kernel'): (1, 768, 768)
	('DiTBlock_0', 'MlpBlock_0', 'Dense_0', 'bias'): (1, 3072)
	('DiTBlock_0', 'MlpBlock_0', 'Dense_0', 'kernel'): (1, 768, 3072)
	('DiTBlock_0', 'MlpBlock_0', 'Dense_1', 'bias'): (1, 768)
	('DiTBlock_0', 'MlpBlock_0', 'Dense_1', 'kernel'): (1, 3072, 768)
	('DiTBlock_1', 'Dense_0', 'bias'): (1, 4608)
	('DiTBlock_1', 'Dense_0', 'kernel'): (1, 768, 4608)
	('DiTBlock_1', 'Dense_1', 'bias'): (1, 768)
	('DiTBlock_1', 'Dense_1', 'kernel'): (1, 768, 768)
	('DiTBlock_1', 'Dense_2', 'bias'): (1, 768)
	('DiTBlock_1', 'Dense_2', 'kernel'): (1, 768, 768)
	('DiTBlock_1', 'Dense_3', 'bias'): (1, 768)
	('DiTBlock_1', 'Dense_3', 'kernel'): (1, 768, 768)
	('DiTBlock_1', 'Dense_4', 'bias'): (1, 768)
	('DiTBlock_1', 'Dense_4', 'kernel'): (1, 768, 768)
	('DiTBlock_1', 'MlpBlock_0', 'Dense_0', 'bias'): (1, 3072)
	('DiTBlock_1', 'MlpBlock_0', 'Dense_0', 'kernel'): (1, 768, 3072)
	('DiTBlock_1', 'MlpBlock_0', 'Dense_1', 'bias'): (1, 768)
	('DiTBlock_1', 'MlpBlock_0', 'Dense_1', 'kernel'): (1, 3072, 768)
	('DiTBlock_10', 'Dense_0', 'bias'): (1, 4608)
	('DiTBlock_10', 'Dense_0', 'kernel'): (1, 768, 4608)
	('DiTBlock_10', 'Dense_1', 'bias'): (1, 768)
	('DiTBlock_10', 'Dense_1', 'kernel'): (1, 768, 768)
	('DiTBlock_10', 'Dense_2', 'bias'): (1, 768)
	('DiTBlock_10', 'Dense_2', 'kernel'): (1, 768, 768)
	('DiTBlock_10', 'Dense_3', 'bias'): (1, 768)
	('DiTBlock_10', 'Dense_3', 'kernel'): (1, 768, 768)
	('DiTBlock_10', 'Dense_4', 'bias'): (1, 768)
	('DiTBlock_10', 'Dense_4', 'kernel'): (1, 768, 768)
	('DiTBlock_10', 'MlpBlock_0', 'Dense_0', 'bias'): (1, 3072)
	('DiTBlock_10', 'MlpBlock_0', 'Dense_0', 'kernel'): (1, 768, 3072)
	('DiTBlock_10', 'MlpBlock_0', 'Dense_1', 'bias'): (1, 768)
	('DiTBlock_10', 'MlpBlock_0', 'Dense_1', 'kernel'): (1, 3072, 768)
	('DiTBlock_11', 'Dense_0', 'bias'): (1, 4608)
	('DiTBlock_11', 'Dense_0', 'kernel'): (1, 768, 4608)
	('DiTBlock_11', 'Dense_1', 'bias'): (1, 768)
	('DiTBlock_11', 'Dense_1', 'kernel'): (1, 768, 768)
	('DiTBlock_11', 'Dense_2', 'bias'): (1, 768)
	('DiTBlock_11', 'Dense_2', 'kernel'): (1, 768, 768)
	('DiTBlock_11', 'Dense_3', 'bias'): (1, 768)
	('DiTBlock_11', 'Dense_3', 'kernel'): (1, 768, 768)
	('DiTBlock_11', 'Dense_4', 'bias'): (1, 768)
	('DiTBlock_11', 'Dense_4', 'kernel'): (1, 768, 768)
	('DiTBlock_11', 'MlpBlock_0', 'Dense_0', 'bias'): (1, 3072)
	('DiTBlock_11', 'MlpBlock_0', 'Dense_0', 'kernel'): (1, 768, 3072)
	('DiTBlock_11', 'MlpBlock_0', 'Dense_1', 'bias'): (1, 768)
	('DiTBlock_11', 'MlpBlock_0', 'Dense_1', 'kernel'): (1, 3072, 768)
	('DiTBlock_2', 'Dense_0', 'bias'): (1, 4608)
	('DiTBlock_2', 'Dense_0', 'kernel'): (1, 768, 4608)
	('DiTBlock_2', 'Dense_1', 'bias'): (1, 768)
	('DiTBlock_2', 'Dense_1', 'kernel'): (1, 768, 768)
	('DiTBlock_2', 'Dense_2', 'bias'): (1, 768)
	('DiTBlock_2', 'Dense_2', 'kernel'): (1, 768, 768)
	('DiTBlock_2', 'Dense_3', 'bias'): (1, 768)
	('DiTBlock_2', 'Dense_3', 'kernel'): (1, 768, 768)
	('DiTBlock_2', 'Dense_4', 'bias'): (1, 768)
	('DiTBlock_2', 'Dense_4', 'kernel'): (1, 768, 768)
	('DiTBlock_2', 'MlpBlock_0', 'Dense_0', 'bias'): (1, 3072)
	('DiTBlock_2', 'MlpBlock_0', 'Dense_0', 'kernel'): (1, 768, 3072)
	('DiTBlock_2', 'MlpBlock_0', 'Dense_1', 'bias'): (1, 768)
	('DiTBlock_2', 'MlpBlock_0', 'Dense_1', 'kernel'): (1, 3072, 768)
	('DiTBlock_3', 'Dense_0', 'bias'): (1, 4608)
	('DiTBlock_3', 'Dense_0', 'kernel'): (1, 768, 4608)
	('DiTBlock_3', 'Dense_1', 'bias'): (1, 768)
	('DiTBlock_3', 'Dense_1', 'kernel'): (1, 768, 768)
	('DiTBlock_3', 'Dense_2', 'bias'): (1, 768)
	('DiTBlock_3', 'Dense_2', 'kernel'): (1, 768, 768)
	('DiTBlock_3', 'Dense_3', 'bias'): (1, 768)
	('DiTBlock_3', 'Dense_3', 'kernel'): (1, 768, 768)
	('DiTBlock_3', 'Dense_4', 'bias'): (1, 768)
	('DiTBlock_3', 'Dense_4', 'kernel'): (1, 768, 768)
	('DiTBlock_3', 'MlpBlock_0', 'Dense_0', 'bias'): (1, 3072)
	('DiTBlock_3', 'MlpBlock_0', 'Dense_0', 'kernel'): (1, 768, 3072)
	('DiTBlock_3', 'MlpBlock_0', 'Dense_1', 'bias'): (1, 768)
	('DiTBlock_3', 'MlpBlock_0', 'Dense_1', 'kernel'): (1, 3072, 768)
	('DiTBlock_4', 'Dense_0', 'bias'): (1, 4608)
	('DiTBlock_4', 'Dense_0', 'kernel'): (1, 768, 4608)
	('DiTBlock_4', 'Dense_1', 'bias'): (1, 768)
	('DiTBlock_4', 'Dense_1', 'kernel'): (1, 768, 768)
	('DiTBlock_4', 'Dense_2', 'bias'): (1, 768)
	('DiTBlock_4', 'Dense_2', 'kernel'): (1, 768, 768)
	('DiTBlock_4', 'Dense_3', 'bias'): (1, 768)
	('DiTBlock_4', 'Dense_3', 'kernel'): (1, 768, 768)
	('DiTBlock_4', 'Dense_4', 'bias'): (1, 768)
	('DiTBlock_4', 'Dense_4', 'kernel'): (1, 768, 768)
	('DiTBlock_4', 'MlpBlock_0', 'Dense_0', 'bias'): (1, 3072)
	('DiTBlock_4', 'MlpBlock_0', 'Dense_0', 'kernel'): (1, 768, 3072)
	('DiTBlock_4', 'MlpBlock_0', 'Dense_1', 'bias'): (1, 768)
	('DiTBlock_4', 'MlpBlock_0', 'Dense_1', 'kernel'): (1, 3072, 768)
	('DiTBlock_5', 'Dense_0', 'bias'): (1, 4608)
	('DiTBlock_5', 'Dense_0', 'kernel'): (1, 768, 4608)
	('DiTBlock_5', 'Dense_1', 'bias'): (1, 768)
	('DiTBlock_5', 'Dense_1', 'kernel'): (1, 768, 768)
	('DiTBlock_5', 'Dense_2', 'bias'): (1, 768)
	('DiTBlock_5', 'Dense_2', 'kernel'): (1, 768, 768)
	('DiTBlock_5', 'Dense_3', 'bias'): (1, 768)
	('DiTBlock_5', 'Dense_3', 'kernel'): (1, 768, 768)
	('DiTBlock_5', 'Dense_4', 'bias'): (1, 768)
	('DiTBlock_5', 'Dense_4', 'kernel'): (1, 768, 768)
	('DiTBlock_5', 'MlpBlock_0', 'Dense_0', 'bias'): (1, 3072)
	('DiTBlock_5', 'MlpBlock_0', 'Dense_0', 'kernel'): (1, 768, 3072)
	('DiTBlock_5', 'MlpBlock_0', 'Dense_1', 'bias'): (1, 768)
	('DiTBlock_5', 'MlpBlock_0', 'Dense_1', 'kernel'): (1, 3072, 768)
	('DiTBlock_6', 'Dense_0', 'bias'): (1, 4608)
	('DiTBlock_6', 'Dense_0', 'kernel'): (1, 768, 4608)
	('DiTBlock_6', 'Dense_1', 'bias'): (1, 768)
	('DiTBlock_6', 'Dense_1', 'kernel'): (1, 768, 768)
	('DiTBlock_6', 'Dense_2', 'bias'): (1, 768)
	('DiTBlock_6', 'Dense_2', 'kernel'): (1, 768, 768)
	('DiTBlock_6', 'Dense_3', 'bias'): (1, 768)
	('DiTBlock_6', 'Dense_3', 'kernel'): (1, 768, 768)
	('DiTBlock_6', 'Dense_4', 'bias'): (1, 768)
	('DiTBlock_6', 'Dense_4', 'kernel'): (1, 768, 768)
	('DiTBlock_6', 'MlpBlock_0', 'Dense_0', 'bias'): (1, 3072)
	('DiTBlock_6', 'MlpBlock_0', 'Dense_0', 'kernel'): (1, 768, 3072)
	('DiTBlock_6', 'MlpBlock_0', 'Dense_1', 'bias'): (1, 768)
	('DiTBlock_6', 'MlpBlock_0', 'Dense_1', 'kernel'): (1, 3072, 768)
	('DiTBlock_7', 'Dense_0', 'bias'): (1, 4608)
	('DiTBlock_7', 'Dense_0', 'kernel'): (1, 768, 4608)
	('DiTBlock_7', 'Dense_1', 'bias'): (1, 768)
	('DiTBlock_7', 'Dense_1', 'kernel'): (1, 768, 768)
	('DiTBlock_7', 'Dense_2', 'bias'): (1, 768)
	('DiTBlock_7', 'Dense_2', 'kernel'): (1, 768, 768)
	('DiTBlock_7', 'Dense_3', 'bias'): (1, 768)
	('DiTBlock_7', 'Dense_3', 'kernel'): (1, 768, 768)
	('DiTBlock_7', 'Dense_4', 'bias'): (1, 768)
	('DiTBlock_7', 'Dense_4', 'kernel'): (1, 768, 768)
	('DiTBlock_7', 'MlpBlock_0', 'Dense_0', 'bias'): (1, 3072)
	('DiTBlock_7', 'MlpBlock_0', 'Dense_0', 'kernel'): (1, 768, 3072)
	('DiTBlock_7', 'MlpBlock_0', 'Dense_1', 'bias'): (1, 768)
	('DiTBlock_7', 'MlpBlock_0', 'Dense_1', 'kernel'): (1, 3072, 768)
	('DiTBlock_8', 'Dense_0', 'bias'): (1, 4608)
	('DiTBlock_8', 'Dense_0', 'kernel'): (1, 768, 4608)
	('DiTBlock_8', 'Dense_1', 'bias'): (1, 768)
	('DiTBlock_8', 'Dense_1', 'kernel'): (1, 768, 768)
	('DiTBlock_8', 'Dense_2', 'bias'): (1, 768)
	('DiTBlock_8', 'Dense_2', 'kernel'): (1, 768, 768)
	('DiTBlock_8', 'Dense_3', 'bias'): (1, 768)
	('DiTBlock_8', 'Dense_3', 'kernel'): (1, 768, 768)
	('DiTBlock_8', 'Dense_4', 'bias'): (1, 768)
	('DiTBlock_8', 'Dense_4', 'kernel'): (1, 768, 768)
	('DiTBlock_8', 'MlpBlock_0', 'Dense_0', 'bias'): (1, 3072)
	('DiTBlock_8', 'MlpBlock_0', 'Dense_0', 'kernel'): (1, 768, 3072)
	('DiTBlock_8', 'MlpBlock_0', 'Dense_1', 'bias'): (1, 768)
	('DiTBlock_8', 'MlpBlock_0', 'Dense_1', 'kernel'): (1, 3072, 768)
	('DiTBlock_9', 'Dense_0', 'bias'): (1, 4608)
	('DiTBlock_9', 'Dense_0', 'kernel'): (1, 768, 4608)
	('DiTBlock_9', 'Dense_1', 'bias'): (1, 768)
	('DiTBlock_9', 'Dense_1', 'kernel'): (1, 768, 768)
	('DiTBlock_9', 'Dense_2', 'bias'): (1, 768)
	('DiTBlock_9', 'Dense_2', 'kernel'): (1, 768, 768)
	('DiTBlock_9', 'Dense_3', 'bias'): (1, 768)
	('DiTBlock_9', 'Dense_3', 'kernel'): (1, 768, 768)
	('DiTBlock_9', 'Dense_4', 'bias'): (1, 768)
	('DiTBlock_9', 'Dense_4', 'kernel'): (1, 768, 768)
	('DiTBlock_9', 'MlpBlock_0', 'Dense_0', 'bias'): (1, 3072)
	('DiTBlock_9', 'MlpBlock_0', 'Dense_0', 'kernel'): (1, 768, 3072)
	('DiTBlock_9', 'MlpBlock_0', 'Dense_1', 'bias'): (1, 768)
	('DiTBlock_9', 'MlpBlock_0', 'Dense_1', 'kernel'): (1, 3072, 768)
	('Embed_0', 'embedding'): (1, 256, 1)
	('FinalLayer_0', 'Dense_0', 'bias'): (1, 1536)
	('FinalLayer_0', 'Dense_0', 'kernel'): (1, 768, 1536)
	('FinalLayer_0', 'Dense_1', 'bias'): (1, 16)
	('FinalLayer_0', 'Dense_1', 'kernel'): (1, 768, 16)
	('LabelEmbedder_0', 'Embed_0', 'embedding'): (1, 1001, 768)
	('PatchEmbed_0', 'Conv_0', 'bias'): (1, 768)
	('PatchEmbed_0', 'Conv_0', 'kernel'): (1, 2, 2, 4, 768)
	('TimestepEmbedder_0', 'Dense_0', 'bias'): (1, 768)
	('TimestepEmbedder_0', 'Dense_0', 'kernel'): (1, 256, 768)
	('TimestepEmbedder_0', 'Dense_1', 'bias'): (1, 768)
	('TimestepEmbedder_0', 'Dense_1', 'kernel'): (1, 768, 768)
	('TimestepEmbedder_1', 'Dense_0', 'bias'): (1, 768)
	('TimestepEmbedder_1', 'Dense_0', 'kernel'): (1, 256, 768)
	('TimestepEmbedder_1', 'Dense_1', 'bias'): (1, 768)
	('TimestepEmbedder_1', 'Dense_1', 'kernel'): (1, 768, 768)

	parameter shapes:
	('DiTBlock_0', 'Dense_0', 'bias'): (4608,)
	('DiTBlock_0', 'Dense_0', 'kernel'): (768, 4608)
	('DiTBlock_0', 'Dense_1', 'bias'): (768,)
	('DiTBlock_0', 'Dense_1', 'kernel'): (768, 768)
	('DiTBlock_0', 'Dense_2', 'bias'): (768,)
	('DiTBlock_0', 'Dense_2', 'kernel'): (768, 768)
	('DiTBlock_0', 'Dense_3', 'bias'): (768,)
	('DiTBlock_0', 'Dense_3', 'kernel'): (768, 768)
	('DiTBlock_0', 'Dense_4', 'bias'): (768,)
	('DiTBlock_0', 'Dense_4', 'kernel'): (768, 768)
	('DiTBlock_0', 'MlpBlock_0', 'Dense_0', 'bias'): (3072,)
	('DiTBlock_0', 'MlpBlock_0', 'Dense_0', 'kernel'): (768, 3072)
	('DiTBlock_0', 'MlpBlock_0', 'Dense_1', 'bias'): (768,)
	('DiTBlock_0', 'MlpBlock_0', 'Dense_1', 'kernel'): (3072, 768)
	('DiTBlock_1', 'Dense_0', 'bias'): (4608,)
	('DiTBlock_1', 'Dense_0', 'kernel'): (768, 4608)
	('DiTBlock_1', 'Dense_1', 'bias'): (768,)
	('DiTBlock_1', 'Dense_1', 'kernel'): (768, 768)
	('DiTBlock_1', 'Dense_2', 'bias'): (768,)
	('DiTBlock_1', 'Dense_2', 'kernel'): (768, 768)
	('DiTBlock_1', 'Dense_3', 'bias'): (768,)
	('DiTBlock_1', 'Dense_3', 'kernel'): (768, 768)
	('DiTBlock_1', 'Dense_4', 'bias'): (768,)
	('DiTBlock_1', 'Dense_4', 'kernel'): (768, 768)
	('DiTBlock_1', 'MlpBlock_0', 'Dense_0', 'bias'): (3072,)
	('DiTBlock_1', 'MlpBlock_0', 'Dense_0', 'kernel'): (768, 3072)
	('DiTBlock_1', 'MlpBlock_0', 'Dense_1', 'bias'): (768,)
	('DiTBlock_1', 'MlpBlock_0', 'Dense_1', 'kernel'): (3072, 768)
	('DiTBlock_10', 'Dense_0', 'bias'): (4608,)
	('DiTBlock_10', 'Dense_0', 'kernel'): (768, 4608)
	('DiTBlock_10', 'Dense_1', 'bias'): (768,)
	('DiTBlock_10', 'Dense_1', 'kernel'): (768, 768)
	('DiTBlock_10', 'Dense_2', 'bias'): (768,)
	('DiTBlock_10', 'Dense_2', 'kernel'): (768, 768)
	('DiTBlock_10', 'Dense_3', 'bias'): (768,)
	('DiTBlock_10', 'Dense_3', 'kernel'): (768, 768)
	('DiTBlock_10', 'Dense_4', 'bias'): (768,)
	('DiTBlock_10', 'Dense_4', 'kernel'): (768, 768)
	('DiTBlock_10', 'MlpBlock_0', 'Dense_0', 'bias'): (3072,)
	('DiTBlock_10', 'MlpBlock_0', 'Dense_0', 'kernel'): (768, 3072)
	('DiTBlock_10', 'MlpBlock_0', 'Dense_1', 'bias'): (768,)
	('DiTBlock_10', 'MlpBlock_0', 'Dense_1', 'kernel'): (3072, 768)
	('DiTBlock_11', 'Dense_0', 'bias'): (4608,)
	('DiTBlock_11', 'Dense_0', 'kernel'): (768, 4608)
	('DiTBlock_11', 'Dense_1', 'bias'): (768,)
	('DiTBlock_11', 'Dense_1', 'kernel'): (768, 768)
	('DiTBlock_11', 'Dense_2', 'bias'): (768,)
	('DiTBlock_11', 'Dense_2', 'kernel'): (768, 768)
	('DiTBlock_11', 'Dense_3', 'bias'): (768,)
	('DiTBlock_11', 'Dense_3', 'kernel'): (768, 768)
	('DiTBlock_11', 'Dense_4', 'bias'): (768,)
	('DiTBlock_11', 'Dense_4', 'kernel'): (768, 768)
	('DiTBlock_11', 'MlpBlock_0', 'Dense_0', 'bias'): (3072,)
	('DiTBlock_11', 'MlpBlock_0', 'Dense_0', 'kernel'): (768, 3072)
	('DiTBlock_11', 'MlpBlock_0', 'Dense_1', 'bias'): (768,)
	('DiTBlock_11', 'MlpBlock_0', 'Dense_1', 'kernel'): (3072, 768)
	('DiTBlock_2', 'Dense_0', 'bias'): (4608,)
	('DiTBlock_2', 'Dense_0', 'kernel'): (768, 4608)
	('DiTBlock_2', 'Dense_1', 'bias'): (768,)
	('DiTBlock_2', 'Dense_1', 'kernel'): (768, 768)
	('DiTBlock_2', 'Dense_2', 'bias'): (768,)
	('DiTBlock_2', 'Dense_2', 'kernel'): (768, 768)
	('DiTBlock_2', 'Dense_3', 'bias'): (768,)
	('DiTBlock_2', 'Dense_3', 'kernel'): (768, 768)
	('DiTBlock_2', 'Dense_4', 'bias'): (768,)
	('DiTBlock_2', 'Dense_4', 'kernel'): (768, 768)
	('DiTBlock_2', 'MlpBlock_0', 'Dense_0', 'bias'): (3072,)
	('DiTBlock_2', 'MlpBlock_0', 'Dense_0', 'kernel'): (768, 3072)
	('DiTBlock_2', 'MlpBlock_0', 'Dense_1', 'bias'): (768,)
	('DiTBlock_2', 'MlpBlock_0', 'Dense_1', 'kernel'): (3072, 768)
	('DiTBlock_3', 'Dense_0', 'bias'): (4608,)
	('DiTBlock_3', 'Dense_0', 'kernel'): (768, 4608)
	('DiTBlock_3', 'Dense_1', 'bias'): (768,)
	('DiTBlock_3', 'Dense_1', 'kernel'): (768, 768)
	('DiTBlock_3', 'Dense_2', 'bias'): (768,)
	('DiTBlock_3', 'Dense_2', 'kernel'): (768, 768)
	('DiTBlock_3', 'Dense_3', 'bias'): (768,)
	('DiTBlock_3', 'Dense_3', 'kernel'): (768, 768)
	('DiTBlock_3', 'Dense_4', 'bias'): (768,)
	('DiTBlock_3', 'Dense_4', 'kernel'): (768, 768)
	('DiTBlock_3', 'MlpBlock_0', 'Dense_0', 'bias'): (3072,)
	('DiTBlock_3', 'MlpBlock_0', 'Dense_0', 'kernel'): (768, 3072)
	('DiTBlock_3', 'MlpBlock_0', 'Dense_1', 'bias'): (768,)
	('DiTBlock_3', 'MlpBlock_0', 'Dense_1', 'kernel'): (3072, 768)
	('DiTBlock_4', 'Dense_0', 'bias'): (4608,)
	('DiTBlock_4', 'Dense_0', 'kernel'): (768, 4608)
	('DiTBlock_4', 'Dense_1', 'bias'): (768,)
	('DiTBlock_4', 'Dense_1', 'kernel'): (768, 768)
	('DiTBlock_4', 'Dense_2', 'bias'): (768,)
	('DiTBlock_4', 'Dense_2', 'kernel'): (768, 768)
	('DiTBlock_4', 'Dense_3', 'bias'): (768,)
	('DiTBlock_4', 'Dense_3', 'kernel'): (768, 768)
	('DiTBlock_4', 'Dense_4', 'bias'): (768,)
	('DiTBlock_4', 'Dense_4', 'kernel'): (768, 768)
	('DiTBlock_4', 'MlpBlock_0', 'Dense_0', 'bias'): (3072,)
	('DiTBlock_4', 'MlpBlock_0', 'Dense_0', 'kernel'): (768, 3072)
	('DiTBlock_4', 'MlpBlock_0', 'Dense_1', 'bias'): (768,)
	('DiTBlock_4', 'MlpBlock_0', 'Dense_1', 'kernel'): (3072, 768)
	('DiTBlock_5', 'Dense_0', 'bias'): (4608,)
	('DiTBlock_5', 'Dense_0', 'kernel'): (768, 4608)
	('DiTBlock_5', 'Dense_1', 'bias'): (768,)
	('DiTBlock_5', 'Dense_1', 'kernel'): (768, 768)
	('DiTBlock_5', 'Dense_2', 'bias'): (768,)
	('DiTBlock_5', 'Dense_2', 'kernel'): (768, 768)
	('DiTBlock_5', 'Dense_3', 'bias'): (768,)
	('DiTBlock_5', 'Dense_3', 'kernel'): (768, 768)
	('DiTBlock_5', 'Dense_4', 'bias'): (768,)
	('DiTBlock_5', 'Dense_4', 'kernel'): (768, 768)
	('DiTBlock_5', 'MlpBlock_0', 'Dense_0', 'bias'): (3072,)
	('DiTBlock_5', 'MlpBlock_0', 'Dense_0', 'kernel'): (768, 3072)
	('DiTBlock_5', 'MlpBlock_0', 'Dense_1', 'bias'): (768,)
	('DiTBlock_5', 'MlpBlock_0', 'Dense_1', 'kernel'): (3072, 768)
	('DiTBlock_6', 'Dense_0', 'bias'): (4608,)
	('DiTBlock_6', 'Dense_0', 'kernel'): (768, 4608)
	('DiTBlock_6', 'Dense_1', 'bias'): (768,)
	('DiTBlock_6', 'Dense_1', 'kernel'): (768, 768)
	('DiTBlock_6', 'Dense_2', 'bias'): (768,)
	('DiTBlock_6', 'Dense_2', 'kernel'): (768, 768)
	('DiTBlock_6', 'Dense_3', 'bias'): (768,)
	('DiTBlock_6', 'Dense_3', 'kernel'): (768, 768)
	('DiTBlock_6', 'Dense_4', 'bias'): (768,)
	('DiTBlock_6', 'Dense_4', 'kernel'): (768, 768)
	('DiTBlock_6', 'MlpBlock_0', 'Dense_0', 'bias'): (3072,)
	('DiTBlock_6', 'MlpBlock_0', 'Dense_0', 'kernel'): (768, 3072)
	('DiTBlock_6', 'MlpBlock_0', 'Dense_1', 'bias'): (768,)
	('DiTBlock_6', 'MlpBlock_0', 'Dense_1', 'kernel'): (3072, 768)
	('DiTBlock_7', 'Dense_0', 'bias'): (4608,)
	('DiTBlock_7', 'Dense_0', 'kernel'): (768, 4608)
	('DiTBlock_7', 'Dense_1', 'bias'): (768,)
	('DiTBlock_7', 'Dense_1', 'kernel'): (768, 768)
	('DiTBlock_7', 'Dense_2', 'bias'): (768,)
	('DiTBlock_7', 'Dense_2', 'kernel'): (768, 768)
	('DiTBlock_7', 'Dense_3', 'bias'): (768,)
	('DiTBlock_7', 'Dense_3', 'kernel'): (768, 768)
	('DiTBlock_7', 'Dense_4', 'bias'): (768,)
	('DiTBlock_7', 'Dense_4', 'kernel'): (768, 768)
	('DiTBlock_7', 'MlpBlock_0', 'Dense_0', 'bias'): (3072,)
	('DiTBlock_7', 'MlpBlock_0', 'Dense_0', 'kernel'): (768, 3072)
	('DiTBlock_7', 'MlpBlock_0', 'Dense_1', 'bias'): (768,)
	('DiTBlock_7', 'MlpBlock_0', 'Dense_1', 'kernel'): (3072, 768)
	('DiTBlock_8', 'Dense_0', 'bias'): (4608,)
	('DiTBlock_8', 'Dense_0', 'kernel'): (768, 4608)
	('DiTBlock_8', 'Dense_1', 'bias'): (768,)
	('DiTBlock_8', 'Dense_1', 'kernel'): (768, 768)
	('DiTBlock_8', 'Dense_2', 'bias'): (768,)
	('DiTBlock_8', 'Dense_2', 'kernel'): (768, 768)
	('DiTBlock_8', 'Dense_3', 'bias'): (768,)
	('DiTBlock_8', 'Dense_3', 'kernel'): (768, 768)
	('DiTBlock_8', 'Dense_4', 'bias'): (768,)
	('DiTBlock_8', 'Dense_4', 'kernel'): (768, 768)
	('DiTBlock_8', 'MlpBlock_0', 'Dense_0', 'bias'): (3072,)
	('DiTBlock_8', 'MlpBlock_0', 'Dense_0', 'kernel'): (768, 3072)
	('DiTBlock_8', 'MlpBlock_0', 'Dense_1', 'bias'): (768,)
	('DiTBlock_8', 'MlpBlock_0', 'Dense_1', 'kernel'): (3072, 768)
	('DiTBlock_9', 'Dense_0', 'bias'): (4608,)
	('DiTBlock_9', 'Dense_0', 'kernel'): (768, 4608)
	('DiTBlock_9', 'Dense_1', 'bias'): (768,)
	('DiTBlock_9', 'Dense_1', 'kernel'): (768, 768)
	('DiTBlock_9', 'Dense_2', 'bias'): (768,)
	('DiTBlock_9', 'Dense_2', 'kernel'): (768, 768)
	('DiTBlock_9', 'Dense_3', 'bias'): (768,)
	('DiTBlock_9', 'Dense_3', 'kernel'): (768, 768)
	('DiTBlock_9', 'Dense_4', 'bias'): (768,)
	('DiTBlock_9', 'Dense_4', 'kernel'): (768, 768)
	('DiTBlock_9', 'MlpBlock_0', 'Dense_0', 'bias'): (3072,)
	('DiTBlock_9', 'MlpBlock_0', 'Dense_0', 'kernel'): (768, 3072)
	('DiTBlock_9', 'MlpBlock_0', 'Dense_1', 'bias'): (768,)
	('DiTBlock_9', 'MlpBlock_0', 'Dense_1', 'kernel'): (3072, 768)
	('Embed_0', 'embedding'): (256, 1)
	('FinalLayer_0', 'Dense_0', 'bias'): (1536,)
	('FinalLayer_0', 'Dense_0', 'kernel'): (768, 1536)
	('FinalLayer_0', 'Dense_1', 'bias'): (16,)
	('FinalLayer_0', 'Dense_1', 'kernel'): (768, 16)
	('LabelEmbedder_0', 'Embed_0', 'embedding'): (1001, 768)
	('PatchEmbed_0', 'Conv_0', 'bias'): (768,)
	('PatchEmbed_0', 'Conv_0', 'kernel'): (2, 2, 4, 768)
	('TimestepEmbedder_0', 'Dense_0', 'bias'): (768,)
	('TimestepEmbedder_0', 'Dense_0', 'kernel'): (256, 768)
	('TimestepEmbedder_0', 'Dense_1', 'bias'): (768,)
	('TimestepEmbedder_0', 'Dense_1', 'kernel'): (768, 768)
	('TimestepEmbedder_1', 'Dense_0', 'bias'): (768,)
	('TimestepEmbedder_1', 'Dense_0', 'kernel'): (256, 768)
	('TimestepEmbedder_1', 'Dense_1', 'bias'): (768,)
	('TimestepEmbedder_1', 'Dense_1', 'kernel'): (768, 768)
	┌────────────────────────────────────────────────┐
	│ │
	│ │
	│ │
	│ │
	│ TPU 0,1,2,3 │
	│ │
	│ │
	│ │
	│ │
	└────────────────────────────────────────────────┘
	┌─────────────────────────────────────────────────────────────────────────┐
	│ │
	│ │
	│ │
	│ │
	│ TPU 0,1,2,3 │
	│ │
	│ │
	│ │
	│ │
	└─────────────────────────────────────────────────────────────────────────┘
	doing the else
	(512, 256, 256, 3)
	encode image shape (128, 256, 256, 3)
	Initializing encoder.
	Incoming encoder shape (128, 256, 256, 3)
	Encoder layer (128, 256, 256, 128)
	doing downsample
	Encoder layer (128, 128, 128, 128)
	doing downsample
	Encoder layer (128, 64, 64, 256)
	doing downsample
	Encoder layer (128, 32, 32, 512)
	Encoder layer (128, 32, 32, 512)
	Encoder layer final (128, 32, 32, 512)
	Encoder layer final (128, 32, 32, 512)
	Final embeddings are size (128, 32, 32, 8)
	After quant (128, 32, 32, 4)
	Calc FID for CFG 1.0 and denoise_timesteps 128
	DiT: Input of shape (512, 32, 32, 4) dtype float32
	DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
	DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
	DiT: Conditioning of shape (512, 768) dtype float32
	z_vectors shape (128, 32, 32, 4)
	Decoder incoming shape (128, 32, 32, 4)
	Decoder input (128, 32, 32, 512)
	Mid Block Decoder layer (128, 32, 32, 512)
	Mid Block Decoder layer (128, 32, 32, 512)
	Decoder layer (128, 64, 64, 512)
	Decoder layer (128, 128, 128, 512)
	Decoder layer (128, 256, 256, 256)
	Decoder layer (128, 256, 256, 128)
	FID is 32.973609924316406
	(512, 256, 256, 3)
	Calc FID for CFG 1.0 and denoise_timesteps 64
	DiT: Input of shape (512, 32, 32, 4) dtype float32
	DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
	DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
	DiT: Conditioning of shape (512, 768) dtype float32
	FID is 17.79468536376953
	(512, 256, 256, 3)
	Calc FID for CFG 1.0 and denoise_timesteps 32
	DiT: Input of shape (512, 32, 32, 4) dtype float32
	DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
	DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
	DiT: Conditioning of shape (512, 768) dtype float32
	FID is 17.702611923217773
	(512, 256, 256, 3)
	Calc FID for CFG 1.0 and denoise_timesteps 16
	DiT: Input of shape (512, 32, 32, 4) dtype float32
	DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
	DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
	DiT: Conditioning of shape (512, 768) dtype float32
	FID is 19.416399002075195
	(512, 256, 256, 3)
	Calc FID for CFG 1.0 and denoise_timesteps 8
	DiT: Input of shape (512, 32, 32, 4) dtype float32
	DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
	DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
	DiT: Conditioning of shape (512, 768) dtype float32
	FID is 20.479732513427734
	(512, 256, 256, 3)
	Calc FID for CFG 1.0 and denoise_timesteps 4
	DiT: Input of shape (512, 32, 32, 4) dtype float32
	DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
	DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
	DiT: Conditioning of shape (512, 768) dtype float32
	FID is 22.193859100341797
	(512, 256, 256, 3)
	Calc FID for CFG 1.0 and denoise_timesteps 2
	DiT: Input of shape (512, 32, 32, 4) dtype float32
	DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
	DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
	DiT: Conditioning of shape (512, 768) dtype float32
	FID is 25.887405395507812
	(512, 256, 256, 3)
	Calc FID for CFG 1.0 and denoise_timesteps 1
	DiT: Input of shape (512, 32, 32, 4) dtype float32
	DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
	DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
	DiT: Conditioning of shape (512, 768) dtype float32
	FID is 31.15492820739746
	(512, 256, 256, 3)
	Calc FID for CFG 1.25 and denoise_timesteps 128
	DiT: Input of shape (512, 32, 32, 4) dtype float32
	DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
	DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
	DiT: Conditioning of shape (512, 768) dtype float32
	FID is 20.413524627685547
	(512, 256, 256, 3)
	Calc FID for CFG 1.25 and denoise_timesteps 64
	DiT: Input of shape (512, 32, 32, 4) dtype float32
	DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
	DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
	DiT: Conditioning of shape (512, 768) dtype float32
	FID is 11.41288948059082
	(512, 256, 256, 3)
	Calc FID for CFG 1.25 and denoise_timesteps 32
	DiT: Input of shape (512, 32, 32, 4) dtype float32
	DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
	DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
	DiT: Conditioning of shape (512, 768) dtype float32
	FID is 11.892237663269043
	(512, 256, 256, 3)
	Calc FID for CFG 1.25 and denoise_timesteps 16
	DiT: Input of shape (512, 32, 32, 4) dtype float32
	DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
	DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
	DiT: Conditioning of shape (512, 768) dtype float32
	FID is 13.14644718170166
	(512, 256, 256, 3)
	Calc FID for CFG 1.25 and denoise_timesteps 8
	DiT: Input of shape (512, 32, 32, 4) dtype float32
	DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
	DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
	DiT: Conditioning of shape (512, 768) dtype float32
	FID is 14.370256423950195
	(512, 256, 256, 3)
	Calc FID for CFG 1.25 and denoise_timesteps 4
	DiT: Input of shape (512, 32, 32, 4) dtype float32
	DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
	DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
	DiT: Conditioning of shape (512, 768) dtype float32
	FID is 16.383480072021484
	(512, 256, 256, 3)
	Calc FID for CFG 1.25 and denoise_timesteps 2
	DiT: Input of shape (512, 32, 32, 4) dtype float32
	DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
	DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
	DiT: Conditioning of shape (512, 768) dtype float32
	FID is 20.53908920288086
	(512, 256, 256, 3)
	Calc FID for CFG 1.25 and denoise_timesteps 1
	DiT: Input of shape (512, 32, 32, 4) dtype float32
	DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
	DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
	DiT: Conditioning of shape (512, 768) dtype float32
	FID is 25.01690673828125
	(512, 256, 256, 3)
	Calc FID for CFG 1.5 and denoise_timesteps 128
	DiT: Input of shape (512, 32, 32, 4) dtype float32
	DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
	DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
	DiT: Conditioning of shape (512, 768) dtype float32
	FID is 13.151246070861816
	(512, 256, 256, 3)
	Calc FID for CFG 1.5 and denoise_timesteps 64
	DiT: Input of shape (512, 32, 32, 4) dtype float32
	DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
	DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
	DiT: Conditioning of shape (512, 768) dtype float32
	FID is 9.616386413574219
	(512, 256, 256, 3)
	Calc FID for CFG 1.5 and denoise_timesteps 32
	DiT: Input of shape (512, 32, 32, 4) dtype float32
	DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
	DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
	DiT: Conditioning of shape (512, 768) dtype float32
	FID is 10.276164054870605
	(512, 256, 256, 3)
	Calc FID for CFG 1.5 and denoise_timesteps 16
	DiT: Input of shape (512, 32, 32, 4) dtype float32
	DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
	DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
	DiT: Conditioning of shape (512, 768) dtype float32
	FID is 11.142467498779297
	(512, 256, 256, 3)
	Calc FID for CFG 1.5 and denoise_timesteps 8
	DiT: Input of shape (512, 32, 32, 4) dtype float32
	DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
	DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
	DiT: Conditioning of shape (512, 768) dtype float32
	FID is 12.243997573852539
	(512, 256, 256, 3)
	Calc FID for CFG 1.5 and denoise_timesteps 4
	DiT: Input of shape (512, 32, 32, 4) dtype float32
	DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
	DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
	DiT: Conditioning of shape (512, 768) dtype float32
	FID is 14.081954956054688
	(512, 256, 256, 3)
	Calc FID for CFG 1.5 and denoise_timesteps 2
	DiT: Input of shape (512, 32, 32, 4) dtype float32
	DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
	DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
	DiT: Conditioning of shape (512, 768) dtype float32
	FID is 18.85321807861328
	(512, 256, 256, 3)
	Calc FID for CFG 1.5 and denoise_timesteps 1
	DiT: Input of shape (512, 32, 32, 4) dtype float32
	DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
	DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
	DiT: Conditioning of shape (512, 768) dtype float32
	FID is 24.270793914794922
	(512, 256, 256, 3)
	Calc FID for CFG 1.75 and denoise_timesteps 128
	DiT: Input of shape (512, 32, 32, 4) dtype float32
	DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
	DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
	DiT: Conditioning of shape (512, 768) dtype float32
	FID is 9.665060043334961
	(512, 256, 256, 3)
	Calc FID for CFG 1.75 and denoise_timesteps 64
	DiT: Input of shape (512, 32, 32, 4) dtype float32
	DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
	DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
	DiT: Conditioning of shape (512, 768) dtype float32
	FID is 9.774864196777344
	(512, 256, 256, 3)
	Calc FID for CFG 1.75 and denoise_timesteps 32
	DiT: Input of shape (512, 32, 32, 4) dtype float32
	DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
	DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
	DiT: Conditioning of shape (512, 768) dtype float32
	FID is 10.444902420043945
	(512, 256, 256, 3)
	Calc FID for CFG 1.75 and denoise_timesteps 16
	DiT: Input of shape (512, 32, 32, 4) dtype float32
	DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
	DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
	DiT: Conditioning of shape (512, 768) dtype float32
	FID is 10.88241958618164
	(512, 256, 256, 3)
	Calc FID for CFG 1.75 and denoise_timesteps 8
	DiT: Input of shape (512, 32, 32, 4) dtype float32
	DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
	DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
	DiT: Conditioning of shape (512, 768) dtype float32
	FID is 11.775745391845703
	(512, 256, 256, 3)
	Calc FID for CFG 1.75 and denoise_timesteps 4
	DiT: Input of shape (512, 32, 32, 4) dtype float32
	DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
	DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
	DiT: Conditioning of shape (512, 768) dtype float32
	FID is 13.461782455444336
	(512, 256, 256, 3)
	Calc FID for CFG 1.75 and denoise_timesteps 2
	DiT: Input of shape (512, 32, 32, 4) dtype float32
	DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
	DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
	DiT: Conditioning of shape (512, 768) dtype float32
	FID is 18.71942901611328
	(512, 256, 256, 3)
	Calc FID for CFG 1.75 and denoise_timesteps 1
	DiT: Input of shape (512, 32, 32, 4) dtype float32
	DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
	DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
	DiT: Conditioning of shape (512, 768) dtype float32
	FID is 25.616065979003906
	(512, 256, 256, 3)
	Calc FID for CFG 2.0 and denoise_timesteps 128
	DiT: Input of shape (512, 32, 32, 4) dtype float32
	DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
	DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
	DiT: Conditioning of shape (512, 768) dtype float32
	FID is 8.354022979736328
	(512, 256, 256, 3)
	Calc FID for CFG 2.0 and denoise_timesteps 64
	DiT: Input of shape (512, 32, 32, 4) dtype float32
	DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
	DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
	DiT: Conditioning of shape (512, 768) dtype float32
	FID is 10.746213912963867
	(512, 256, 256, 3)
	Calc FID for CFG 2.0 and denoise_timesteps 32
	DiT: Input of shape (512, 32, 32, 4) dtype float32
	DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
	DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
	DiT: Conditioning of shape (512, 768) dtype float32
	FID is 11.383520126342773
	(512, 256, 256, 3)
	Calc FID for CFG 2.0 and denoise_timesteps 16
	DiT: Input of shape (512, 32, 32, 4) dtype float32
	DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
	DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
	DiT: Conditioning of shape (512, 768) dtype float32
	FID is 11.429499626159668
	(512, 256, 256, 3)
	Calc FID for CFG 2.0 and denoise_timesteps 8
	DiT: Input of shape (512, 32, 32, 4) dtype float32
	DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
	DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
	DiT: Conditioning of shape (512, 768) dtype float32
	FID is 12.006962776184082
	(512, 256, 256, 3)
	Calc FID for CFG 2.0 and denoise_timesteps 4
	DiT: Input of shape (512, 32, 32, 4) dtype float32
	DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
	DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
	DiT: Conditioning of shape (512, 768) dtype float32
	FID is 13.507768630981445
	(512, 256, 256, 3)
	Calc FID for CFG 2.0 and denoise_timesteps 2
	DiT: Input of shape (512, 32, 32, 4) dtype float32
	DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
	DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
	DiT: Conditioning of shape (512, 768) dtype float32
	FID is 19.345657348632812
	(512, 256, 256, 3)
	Calc FID for CFG 2.0 and denoise_timesteps 1
	DiT: Input of shape (512, 32, 32, 4) dtype float32
	DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
	DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
	DiT: Conditioning of shape (512, 768) dtype float32
	FID is 28.042903900146484
	(512, 256, 256, 3)
	Calc FID for CFG 2.25 and denoise_timesteps 128
	DiT: Input of shape (512, 32, 32, 4) dtype float32
	DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
	DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
	DiT: Conditioning of shape (512, 768) dtype float32
	FID is 8.302826881408691
	(512, 256, 256, 3)
	Calc FID for CFG 2.25 and denoise_timesteps 64
	DiT: Input of shape (512, 32, 32, 4) dtype float32
	DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
	DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
	DiT: Conditioning of shape (512, 768) dtype float32
	FID is 11.958938598632812
	(512, 256, 256, 3)
	Calc FID for CFG 2.25 and denoise_timesteps 32
	DiT: Input of shape (512, 32, 32, 4) dtype float32
	DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
	DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
	DiT: Conditioning of shape (512, 768) dtype float32
	FID is 12.487675666809082
	(512, 256, 256, 3)
	Calc FID for CFG 2.25 and denoise_timesteps 16
	DiT: Input of shape (512, 32, 32, 4) dtype float32
	DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
	DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
	DiT: Conditioning of shape (512, 768) dtype float32
	FID is 12.307321548461914
	(512, 256, 256, 3)
	Calc FID for CFG 2.25 and denoise_timesteps 8
	DiT: Input of shape (512, 32, 32, 4) dtype float32
	DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
	DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
	DiT: Conditioning of shape (512, 768) dtype float32
	FID is 12.590841293334961
	(512, 256, 256, 3)
	Calc FID for CFG 2.25 and denoise_timesteps 4
	DiT: Input of shape (512, 32, 32, 4) dtype float32
	DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
	DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
	DiT: Conditioning of shape (512, 768) dtype float32
	FID is 13.891820907592773
	(512, 256, 256, 3)
	Calc FID for CFG 2.25 and denoise_timesteps 2
	DiT: Input of shape (512, 32, 32, 4) dtype float32
	DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
	DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
	DiT: Conditioning of shape (512, 768) dtype float32
	FID is 20.391244888305664
	(512, 256, 256, 3)
	Calc FID for CFG 2.25 and denoise_timesteps 1
	DiT: Input of shape (512, 32, 32, 4) dtype float32
	DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
	DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
	DiT: Conditioning of shape (512, 768) dtype float32
	FID is 31.200162887573242
	(512, 256, 256, 3)
	Calc FID for CFG 2.5 and denoise_timesteps 128
	DiT: Input of shape (512, 32, 32, 4) dtype float32
	DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
	DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
	DiT: Conditioning of shape (512, 768) dtype float32
	FID is 8.865442276000977
	(512, 256, 256, 3)
	Calc FID for CFG 2.5 and denoise_timesteps 64
	DiT: Input of shape (512, 32, 32, 4) dtype float32
	DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
	DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
	DiT: Conditioning of shape (512, 768) dtype float32
	FID is 13.25680923461914
	(512, 256, 256, 3)
	Calc FID for CFG 2.5 and denoise_timesteps 32
	DiT: Input of shape (512, 32, 32, 4) dtype float32
	DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
	DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
	DiT: Conditioning of shape (512, 768) dtype float32
	FID is 13.67801284790039
	(512, 256, 256, 3)
	Calc FID for CFG 2.5 and denoise_timesteps 16
	DiT: Input of shape (512, 32, 32, 4) dtype float32
	DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
	DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
	DiT: Conditioning of shape (512, 768) dtype float32
	FID is 13.289108276367188
	(512, 256, 256, 3)
	Calc FID for CFG 2.5 and denoise_timesteps 8
	DiT: Input of shape (512, 32, 32, 4) dtype float32
	DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
	DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
	DiT: Conditioning of shape (512, 768) dtype float32
	FID is 13.312899589538574
	(512, 256, 256, 3)
	Calc FID for CFG 2.5 and denoise_timesteps 4
	DiT: Input of shape (512, 32, 32, 4) dtype float32
	DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
	DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
	DiT: Conditioning of shape (512, 768) dtype float32
	FID is 14.423715591430664
	(512, 256, 256, 3)
	Calc FID for CFG 2.5 and denoise_timesteps 2
	DiT: Input of shape (512, 32, 32, 4) dtype float32
	DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
	DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
	DiT: Conditioning of shape (512, 768) dtype float32
	FID is 21.74298858642578
	(512, 256, 256, 3)
	Calc FID for CFG 2.5 and denoise_timesteps 1
	DiT: Input of shape (512, 32, 32, 4) dtype float32
	DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
	DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
	DiT: Conditioning of shape (512, 768) dtype float32
	FID is 34.95703125
	(512, 256, 256, 3)
	Calc FID for CFG 2.75 and denoise_timesteps 128
	DiT: Input of shape (512, 32, 32, 4) dtype float32
	DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
	DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
	DiT: Conditioning of shape (512, 768) dtype float32
	FID is 9.744138717651367
	(512, 256, 256, 3)
	Calc FID for CFG 2.75 and denoise_timesteps 64
	DiT: Input of shape (512, 32, 32, 4) dtype float32
	DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
	DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
	DiT: Conditioning of shape (512, 768) dtype float32
	FID is 14.492772102355957
	(512, 256, 256, 3)
	Calc FID for CFG 2.75 and denoise_timesteps 32
	DiT: Input of shape (512, 32, 32, 4) dtype float32
	DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
	DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
	DiT: Conditioning of shape (512, 768) dtype float32
	FID is 14.805295944213867
	(512, 256, 256, 3)
	Calc FID for CFG 2.75 and denoise_timesteps 16
	DiT: Input of shape (512, 32, 32, 4) dtype float32
	DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
	DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
	DiT: Conditioning of shape (512, 768) dtype float32
	FID is 14.210210800170898
	(512, 256, 256, 3)
	Calc FID for CFG 2.75 and denoise_timesteps 8
	DiT: Input of shape (512, 32, 32, 4) dtype float32
	DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
	DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
	DiT: Conditioning of shape (512, 768) dtype float32
	FID is 14.057098388671875
	(512, 256, 256, 3)
	Calc FID for CFG 2.75 and denoise_timesteps 4
	DiT: Input of shape (512, 32, 32, 4) dtype float32
	DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
	DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
	DiT: Conditioning of shape (512, 768) dtype float32
	FID is 15.049466133117676
	(512, 256, 256, 3)
	Calc FID for CFG 2.75 and denoise_timesteps 2
	DiT: Input of shape (512, 32, 32, 4) dtype float32
	DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
	DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
	DiT: Conditioning of shape (512, 768) dtype float32
	FID is 23.29198455810547
	(512, 256, 256, 3)
	Calc FID for CFG 2.75 and denoise_timesteps 1
	DiT: Input of shape (512, 32, 32, 4) dtype float32
	DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
	DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
	DiT: Conditioning of shape (512, 768) dtype float32
	FID is 39.218223571777344
	(512, 256, 256, 3)
	Calc FID for CFG 3.0 and denoise_timesteps 128
	DiT: Input of shape (512, 32, 32, 4) dtype float32
	DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
	DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
	DiT: Conditioning of shape (512, 768) dtype float32
	FID is 10.747198104858398
	(512, 256, 256, 3)
	Calc FID for CFG 3.0 and denoise_timesteps 64
	DiT: Input of shape (512, 32, 32, 4) dtype float32
	DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
	DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
	DiT: Conditioning of shape (512, 768) dtype float32
	FID is 15.59070873260498
	(512, 256, 256, 3)
	Calc FID for CFG 3.0 and denoise_timesteps 32
	DiT: Input of shape (512, 32, 32, 4) dtype float32
	DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
	DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
	DiT: Conditioning of shape (512, 768) dtype float32
	FID is 15.845895767211914
	(512, 256, 256, 3)
	Calc FID for CFG 3.0 and denoise_timesteps 16
	DiT: Input of shape (512, 32, 32, 4) dtype float32
	DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
	DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
	DiT: Conditioning of shape (512, 768) dtype float32
	FID is 15.133733749389648
	(512, 256, 256, 3)
	Calc FID for CFG 3.0 and denoise_timesteps 8
	DiT: Input of shape (512, 32, 32, 4) dtype float32
	DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
	DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
	DiT: Conditioning of shape (512, 768) dtype float32
	FID is 14.794069290161133
	(512, 256, 256, 3)
	Calc FID for CFG 3.0 and denoise_timesteps 4
	DiT: Input of shape (512, 32, 32, 4) dtype float32
	DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
	DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
	DiT: Conditioning of shape (512, 768) dtype float32
	FID is 15.722245216369629
	(512, 256, 256, 3)
	Calc FID for CFG 3.0 and denoise_timesteps 2
	DiT: Input of shape (512, 32, 32, 4) dtype float32
	DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
	DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
	DiT: Conditioning of shape (512, 768) dtype float32
	FID is 25.04941749572754
	(512, 256, 256, 3)
	Calc FID for CFG 3.0 and denoise_timesteps 1
	DiT: Input of shape (512, 32, 32, 4) dtype float32
	DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
	DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
	DiT: Conditioning of shape (512, 768) dtype float32
	FID is 43.867088317871094