Upload folder using huggingface_hub

4024e46 verified 4 months ago

212 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
	/tmp/tmplnfzqxrg
	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: naive
	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 10646780 seconds ago.
	Loaded model with step 960001
	┌──────────────────────────────────────────────────────────────────────────────┐
	│ 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 10105819 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.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
	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 12.313458442687988
	(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 8.551166534423828
	(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 7.306366920471191
	(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 7.4663262367248535
	(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.268986701965332
	[1;34mwandb[0m:
	[1;34mwandb[0m: 🚀 View run [33mshortcut_imagenet256[0m at: [34mhttps://wandb.ai/daniel-z-kaplan/shortcut/runs/shortcut_imagenet256_20260103_232726_345353_10[0m
	[1;34mwandb[0m: Find logs at: [1;35m../../../tmp/tmplnfzqxrg/wandb/run-20260103_232727-shortcut_imagenet256_20260103_232726_345353_10/logs[0m