diff --git "a/ae_sym_4ch_large/logEma.txt" "b/ae_sym_4ch_large/logEma.txt" new file mode 100644--- /dev/null +++ "b/ae_sym_4ch_large/logEma.txt" @@ -0,0 +1,2254 @@ +Using devices [TpuDevice(id=0, process_index=0, coords=(0,0,0), core_on_chip=0), TpuDevice(id=1, process_index=0, coords=(1,0,0), core_on_chip=0), TpuDevice(id=2, process_index=0, coords=(0,1,0), core_on_chip=0), TpuDevice(id=3, process_index=0, coords=(1,1,0), core_on_chip=0)] +Device count 4 +Global device count 4 +Global Batch: 512 +Node Batch: 512 +Device Batch: 128 +Loading dataset +Loading dataset +creating model +beta1: 0.9 +beta2: 0.999 +bootstrap_cfg: 1 +bootstrap_dt_bias: 0 +bootstrap_ema: 1 +bootstrap_every: 8 +cfg_scale: 1.5 +class_dropout_prob: 0.1 +denoise_timesteps: 128 +depth: 24 +dropout: 0.0 +dt_sampling: uniform +hidden_size: 1024 +lr: 0.0001 +mlp_ratio: 4 +num_classes: 1000 +num_heads: 16 +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, 4) +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: 67563271 +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 17220733 seconds ago. +Loaded model with step 511001 +┌──────────────────────────────────────────────────────────────────────────────┐ +│ 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, 4) +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, 1024) dtype bfloat16 +DiT: Patch Embed of shape (4, 256, 1024) dtype bfloat16 +DiT: Conditioning of shape (1, 1024) dtype float32 + + DiT Summary  +┏━━━━━━━━━━━━━━━━━━━━━━━━━━��━━━━━━━┳━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓ +┃ path  ┃ module  ┃ inputs  ┃ outputs  ┃ params  ┃ +┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┩ +│ │ DiT │ - float32[4,32,32,4] │ bfloat16[4,32,32,4] │ │ +│ │ │ - float32[1] │ │ │ +│ │ │ - float32[1] │ │ │ +│ │ │ - int32[1] │ │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ PatchEmbed_0 │ PatchEmbed │ float32[4,32,32,4] │ bfloat16[4,256,1024] │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ PatchEmbed_0/Conv_0 │ Conv │ float32[4,32,32,4] │ bfloat16[4,16,16,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[2,2,4,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 17,408 (69.6 KB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ TimestepEmbedder_0 │ TimestepEmbedder │ float32[1] │ float32[1,1024] │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ TimestepEmbedder_0/Dense_0 │ Dense │ bfloat16[1,256] │ bfloat16[1,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[256,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 263,168 (1.1 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ TimestepEmbedder_0/Dense_1 │ Dense │ bfloat16[1,1024] │ float32[1,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ TimestepEmbedder_1 │ TimestepEmbedder │ float32[1] │ float32[1,1024] │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ TimestepEmbedder_1/Dense_0 │ Dense │ bfloat16[1,256] │ bfloat16[1,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[256,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 263,168 (1.1 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ TimestepEmbedder_1/Dense_1 │ Dense │ bfloat16[1,1024] │ float32[1,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ LabelEmbedder_0 │ LabelEmbedder │ int32[1] │ bfloat16[1,1024] │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ LabelEmbedder_0/Embed_0 │ Embed │ int32[1] │ bfloat16[1,1024] │ embedding: float32[1001,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,025,024 (4.1 MB) │ +├─────���────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_0 │ DiTBlock │ - bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ │ +│ │ │ - float32[1,1024] │ │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_0/Dense_0 │ Dense │ float32[1,1024] │ bfloat16[1,6144] │ bias: float32[6144] │ +│ │ │ │ │ kernel: float32[1024,6144] │ +│ │ │ │ │ │ +│ │ │ │ │ 6,297,600 (25.2 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_0/LayerNorm_0 │ LayerNorm │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_0/Dense_1 │ Dense │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_0/Dense_2 │ Dense │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_0/Dense_3 │ Dense │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_0/Dense_4 │ Dense │ float32[4,256,1024] │ bfloat16[4,256,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_0/LayerNorm_1 │ LayerNorm │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_0/MlpBlock_0 │ MlpBlock │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_0/MlpBlock_0/Dense_0 │ Dense │ bfloat16[4,256,1024] │ bfloat16[4,256,4096] │ bias: float32[4096] │ +│ │ │ │ │ kernel: float32[1024,4096] │ +│ │ │ │ │ │ +│ │ │ │ │ 4,198,400 (16.8 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_0/MlpBlock_0/Dropout_0 │ Dropout │ bfloat16[4,256,4096] │ bfloat16[4,256,4096] │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_0/MlpBlock_0/Dense_1 │ Dense │ bfloat16[4,256,4096] │ bfloat16[4,256,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[4096,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 4,195,328 (16.8 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_0/MlpBlock_0/Dropout_1 │ Dropout │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_1 │ DiTBlock │ - bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ │ +│ │ │ - float32[1,1024] │ │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_1/Dense_0 │ Dense │ float32[1,1024] │ bfloat16[1,6144] │ bias: float32[6144] │ +│ │ │ │ │ kernel: float32[1024,6144] │ +│ │ │ │ │ │ +│ │ │ │ │ 6,297,600 (25.2 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_1/LayerNorm_0 │ LayerNorm │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_1/Dense_1 │ Dense │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_1/Dense_2 │ Dense │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_1/Dense_3 │ Dense │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_1/Dense_4 │ Dense │ float32[4,256,1024] │ bfloat16[4,256,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_1/LayerNorm_1 │ LayerNorm │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_1/MlpBlock_0 │ MlpBlock │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_1/MlpBlock_0/Dense_0 │ Dense │ bfloat16[4,256,1024] │ bfloat16[4,256,4096] │ bias: float32[4096] │ +│ │ │ │ │ kernel: float32[1024,4096] │ +│ │ │ │ │ │ +│ │ │ │ │ 4,198,400 (16.8 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_1/MlpBlock_0/Dropout_0 │ Dropout │ bfloat16[4,256,4096] │ bfloat16[4,256,4096] │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_1/MlpBlock_0/Dense_1 │ Dense │ bfloat16[4,256,4096] │ bfloat16[4,256,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[4096,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 4,195,328 (16.8 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_1/MlpBlock_0/Dropout_1 │ Dropout │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_2 │ DiTBlock │ - bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ │ +│ │ │ - float32[1,1024] │ │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_2/Dense_0 │ Dense │ float32[1,1024] │ bfloat16[1,6144] │ bias: float32[6144] │ +│ │ │ │ │ kernel: float32[1024,6144] │ +│ │ │ │ │ │ +│ │ │ │ │ 6,297,600 (25.2 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_2/LayerNorm_0 │ LayerNorm │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_2/Dense_1 │ Dense │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_2/Dense_2 │ Dense │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_2/Dense_3 │ Dense │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_2/Dense_4 │ Dense │ float32[4,256,1024] │ bfloat16[4,256,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_2/LayerNorm_1 │ LayerNorm │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_2/MlpBlock_0 │ MlpBlock │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_2/MlpBlock_0/Dense_0 │ Dense │ bfloat16[4,256,1024] │ bfloat16[4,256,4096] │ bias: float32[4096] │ +│ │ │ │ │ kernel: float32[1024,4096] │ +│ │ │ │ │ │ +│ │ │ │ │ 4,198,400 (16.8 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_2/MlpBlock_0/Dropout_0 │ Dropout │ bfloat16[4,256,4096] │ bfloat16[4,256,4096] │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_2/MlpBlock_0/Dense_1 │ Dense │ bfloat16[4,256,4096] │ bfloat16[4,256,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[4096,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 4,195,328 (16.8 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_2/MlpBlock_0/Dropout_1 │ Dropout │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_3 │ DiTBlock │ - bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ │ +│ │ │ - float32[1,1024] │ │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_3/Dense_0 │ Dense │ float32[1,1024] │ bfloat16[1,6144] │ bias: float32[6144] │ +│ │ │ │ │ kernel: float32[1024,6144] │ +│ │ │ │ │ │ +│ │ │ │ │ 6,297,600 (25.2 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_3/LayerNorm_0 │ LayerNorm │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ │ +├──────────────────────────────────┼─────────────��────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_3/Dense_1 │ Dense │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_3/Dense_2 │ Dense │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_3/Dense_3 │ Dense │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_3/Dense_4 │ Dense │ float32[4,256,1024] │ bfloat16[4,256,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_3/LayerNorm_1 │ LayerNorm │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_3/MlpBlock_0 │ MlpBlock │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_3/MlpBlock_0/Dense_0 │ Dense │ bfloat16[4,256,1024] │ bfloat16[4,256,4096] │ bias: float32[4096] │ +│ │ │ │ │ kernel: float32[1024,4096] │ +│ │ │ │ │ │ +│ │ │ │ │ 4,198,400 (16.8 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_3/MlpBlock_0/Dropout_0 │ Dropout │ bfloat16[4,256,4096] │ bfloat16[4,256,4096] │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_3/MlpBlock_0/Dense_1 │ Dense │ bfloat16[4,256,4096] │ bfloat16[4,256,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[4096,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 4,195,328 (16.8 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_3/MlpBlock_0/Dropout_1 │ Dropout │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_4 │ DiTBlock │ - bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ │ +│ │ │ - float32[1,1024] │ │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_4/Dense_0 │ Dense │ float32[1,1024] │ bfloat16[1,6144] │ bias: float32[6144] │ +│ │ │ │ │ kernel: float32[1024,6144] │ +│ │ │ │ │ │ +│ │ │ │ │ 6,297,600 (25.2 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_4/LayerNorm_0 │ LayerNorm │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_4/Dense_1 │ Dense │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_4/Dense_2 │ Dense │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_4/Dense_3 │ Dense │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_4/Dense_4 │ Dense │ float32[4,256,1024] │ bfloat16[4,256,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼────────���─────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_4/LayerNorm_1 │ LayerNorm │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_4/MlpBlock_0 │ MlpBlock │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_4/MlpBlock_0/Dense_0 │ Dense │ bfloat16[4,256,1024] │ bfloat16[4,256,4096] │ bias: float32[4096] │ +│ │ │ │ │ kernel: float32[1024,4096] │ +│ │ │ │ │ │ +│ │ │ │ │ 4,198,400 (16.8 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_4/MlpBlock_0/Dropout_0 │ Dropout │ bfloat16[4,256,4096] │ bfloat16[4,256,4096] │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_4/MlpBlock_0/Dense_1 │ Dense │ bfloat16[4,256,4096] │ bfloat16[4,256,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[4096,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 4,195,328 (16.8 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_4/MlpBlock_0/Dropout_1 │ Dropout │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_5 │ DiTBlock │ - bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ │ +│ │ │ - float32[1,1024] │ │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_5/Dense_0 │ Dense │ float32[1,1024] │ bfloat16[1,6144] │ bias: float32[6144] │ +│ │ │ │ │ kernel: float32[1024,6144] │ +│ │ │ │ │ │ +│ │ │ │ │ 6,297,600 (25.2 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_5/LayerNorm_0 │ LayerNorm │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_5/Dense_1 │ Dense │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_5/Dense_2 │ Dense │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_5/Dense_3 │ Dense │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼─────��─────────────────────────┤ +│ DiTBlock_5/Dense_4 │ Dense │ float32[4,256,1024] │ bfloat16[4,256,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_5/LayerNorm_1 │ LayerNorm │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_5/MlpBlock_0 │ MlpBlock │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_5/MlpBlock_0/Dense_0 │ Dense │ bfloat16[4,256,1024] │ bfloat16[4,256,4096] │ bias: float32[4096] │ +│ │ │ │ │ kernel: float32[1024,4096] │ +│ │ │ │ │ │ +│ │ │ │ │ 4,198,400 (16.8 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_5/MlpBlock_0/Dropout_0 │ Dropout │ bfloat16[4,256,4096] │ bfloat16[4,256,4096] │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_5/MlpBlock_0/Dense_1 │ Dense │ bfloat16[4,256,4096] │ bfloat16[4,256,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[4096,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 4,195,328 (16.8 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_5/MlpBlock_0/Dropout_1 │ Dropout │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_6 │ DiTBlock │ - bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ │ +│ │ │ - float32[1,1024] │ │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_6/Dense_0 │ Dense │ float32[1,1024] │ bfloat16[1,6144] │ bias: float32[6144] │ +│ │ │ │ │ kernel: float32[1024,6144] │ +│ │ │ │ │ │ +│ │ │ │ │ 6,297,600 (25.2 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_6/LayerNorm_0 │ LayerNorm │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_6/Dense_1 │ Dense │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_6/Dense_2 │ Dense │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_6/Dense_3 │ Dense │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_6/Dense_4 │ Dense │ float32[4,256,1024] │ bfloat16[4,256,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_6/LayerNorm_1 │ LayerNorm │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_6/MlpBlock_0 │ MlpBlock │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_6/MlpBlock_0/Dense_0 │ Dense │ bfloat16[4,256,1024] │ bfloat16[4,256,4096] │ bias: float32[4096] │ +│ │ │ │ │ kernel: float32[1024,4096] │ +│ │ │ │ │ │ +│ │ │ │ │ 4,198,400 (16.8 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_6/MlpBlock_0/Dropout_0 │ Dropout │ bfloat16[4,256,4096] │ bfloat16[4,256,4096] │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_6/MlpBlock_0/Dense_1 │ Dense │ bfloat16[4,256,4096] │ bfloat16[4,256,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[4096,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 4,195,328 (16.8 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_6/MlpBlock_0/Dropout_1 │ Dropout │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_7 │ DiTBlock │ - bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ │ +│ │ │ - float32[1,1024] │ │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_7/Dense_0 │ Dense │ float32[1,1024] │ bfloat16[1,6144] │ bias: float32[6144] │ +│ │ │ │ │ kernel: float32[1024,6144] │ +│ │ │ │ │ │ +│ │ │ │ │ 6,297,600 (25.2 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_7/LayerNorm_0 │ LayerNorm │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_7/Dense_1 │ Dense │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_7/Dense_2 │ Dense │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_7/Dense_3 │ Dense │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_7/Dense_4 │ Dense │ float32[4,256,1024] │ bfloat16[4,256,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_7/LayerNorm_1 │ LayerNorm │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_7/MlpBlock_0 │ MlpBlock │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_7/MlpBlock_0/Dense_0 │ Dense │ bfloat16[4,256,1024] │ bfloat16[4,256,4096] │ bias: float32[4096] │ +│ │ │ │ │ kernel: float32[1024,4096] │ +│ │ │ │ │ │ +│ │ │ │ │ 4,198,400 (16.8 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_7/MlpBlock_0/Dropout_0 │ Dropout │ bfloat16[4,256,4096] │ bfloat16[4,256,4096] │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_7/MlpBlock_0/Dense_1 │ Dense │ bfloat16[4,256,4096] │ bfloat16[4,256,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[4096,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 4,195,328 (16.8 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_7/MlpBlock_0/Dropout_1 │ Dropout │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_8 │ DiTBlock │ - bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ │ +│ │ │ - float32[1,1024] │ │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_8/Dense_0 │ Dense │ float32[1,1024] │ bfloat16[1,6144] │ bias: float32[6144] │ +│ │ │ │ │ kernel: float32[1024,6144] │ +│ │ │ │ │ │ +│ │ │ │ │ 6,297,600 (25.2 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_8/LayerNorm_0 │ LayerNorm │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_8/Dense_1 │ Dense │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_8/Dense_2 │ Dense │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_8/Dense_3 │ Dense │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_8/Dense_4 │ Dense │ float32[4,256,1024] │ bfloat16[4,256,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_8/LayerNorm_1 │ LayerNorm │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_8/MlpBlock_0 │ MlpBlock │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_8/MlpBlock_0/Dense_0 │ Dense │ bfloat16[4,256,1024] │ bfloat16[4,256,4096] │ bias: float32[4096] │ +│ │ │ │ │ kernel: float32[1024,4096] │ +│ │ │ │ │ │ +│ │ │ │ │ 4,198,400 (16.8 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_8/MlpBlock_0/Dropout_0 │ Dropout │ bfloat16[4,256,4096] │ bfloat16[4,256,4096] │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_8/MlpBlock_0/Dense_1 │ Dense │ bfloat16[4,256,4096] │ bfloat16[4,256,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[4096,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 4,195,328 (16.8 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_8/MlpBlock_0/Dropout_1 │ Dropout │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_9 │ DiTBlock │ - bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ │ +│ │ │ - float32[1,1024] │ │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_9/Dense_0 │ Dense │ float32[1,1024] │ bfloat16[1,6144] │ bias: float32[6144] │ +│ │ │ │ │ kernel: float32[1024,6144] │ +│ │ │ │ │ │ +│ │ │ │ │ 6,297,600 (25.2 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_9/LayerNorm_0 │ LayerNorm │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ │ +├───────────────────���──────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_9/Dense_1 │ Dense │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_9/Dense_2 │ Dense │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_9/Dense_3 │ Dense │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_9/Dense_4 │ Dense │ float32[4,256,1024] │ bfloat16[4,256,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_9/LayerNorm_1 │ LayerNorm │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_9/MlpBlock_0 │ MlpBlock │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_9/MlpBlock_0/Dense_0 │ Dense │ bfloat16[4,256,1024] │ bfloat16[4,256,4096] │ bias: float32[4096] │ +│ │ │ │ │ kernel: float32[1024,4096] │ +│ │ │ │ │ │ +│ │ │ │ │ 4,198,400 (16.8 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_9/MlpBlock_0/Dropout_0 │ Dropout │ bfloat16[4,256,4096] │ bfloat16[4,256,4096] │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_9/MlpBlock_0/Dense_1 │ Dense │ bfloat16[4,256,4096] │ bfloat16[4,256,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[4096,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 4,195,328 (16.8 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_9/MlpBlock_0/Dropout_1 │ Dropout │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_10 │ DiTBlock │ - bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ │ +│ │ │ - float32[1,1024] │ │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_10/Dense_0 │ Dense │ float32[1,1024] │ bfloat16[1,6144] │ bias: float32[6144] │ +│ │ │ │ │ kernel: float32[1024,6144] │ +│ │ │ │ │ │ +│ │ │ │ │ 6,297,600 (25.2 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_10/LayerNorm_0 │ LayerNorm │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_10/Dense_1 │ Dense │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_10/Dense_2 │ Dense │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_10/Dense_3 │ Dense │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_10/Dense_4 │ Dense │ float32[4,256,1024] │ bfloat16[4,256,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_10/LayerNorm_1 │ LayerNorm │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_10/MlpBlock_0 │ MlpBlock │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_10/MlpBlock_0/Dense_0 │ Dense │ bfloat16[4,256,1024] │ bfloat16[4,256,4096] │ bias: float32[4096] │ +│ │ │ │ │ kernel: float32[1024,4096] │ +│ │ │ │ │ │ +│ │ │ │ │ 4,198,400 (16.8 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_10/MlpBlock_0/Dropout_0 │ Dropout │ bfloat16[4,256,4096] │ bfloat16[4,256,4096] │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_10/MlpBlock_0/Dense_1 │ Dense │ bfloat16[4,256,4096] │ bfloat16[4,256,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[4096,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 4,195,328 (16.8 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_10/MlpBlock_0/Dropout_1 │ Dropout │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_11 │ DiTBlock │ - bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ │ +│ │ │ - float32[1,1024] │ │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_11/Dense_0 │ Dense │ float32[1,1024] │ bfloat16[1,6144] │ bias: float32[6144] │ +│ │ │ │ │ kernel: float32[1024,6144] │ +│ │ │ │ │ │ +│ │ │ │ │ 6,297,600 (25.2 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_11/LayerNorm_0 │ LayerNorm │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_11/Dense_1 │ Dense │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_11/Dense_2 │ Dense │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_11/Dense_3 │ Dense │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼─���──────────────────────┼───────────────────────────────┤ +│ DiTBlock_11/Dense_4 │ Dense │ float32[4,256,1024] │ bfloat16[4,256,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_11/LayerNorm_1 │ LayerNorm │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_11/MlpBlock_0 │ MlpBlock │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_11/MlpBlock_0/Dense_0 │ Dense │ bfloat16[4,256,1024] │ bfloat16[4,256,4096] │ bias: float32[4096] │ +│ │ │ │ │ kernel: float32[1024,4096] │ +│ │ │ │ │ │ +│ │ │ │ │ 4,198,400 (16.8 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_11/MlpBlock_0/Dropout_0 │ Dropout │ bfloat16[4,256,4096] │ bfloat16[4,256,4096] │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_11/MlpBlock_0/Dense_1 │ Dense │ bfloat16[4,256,4096] │ bfloat16[4,256,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[4096,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 4,195,328 (16.8 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_11/MlpBlock_0/Dropout_1 │ Dropout │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_12 │ DiTBlock │ - bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ │ +│ │ │ - float32[1,1024] │ │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_12/Dense_0 │ Dense │ float32[1,1024] │ bfloat16[1,6144] │ bias: float32[6144] │ +│ │ │ │ │ kernel: float32[1024,6144] │ +│ │ │ │ │ │ +│ │ │ │ │ 6,297,600 (25.2 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_12/LayerNorm_0 │ LayerNorm │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_12/Dense_1 │ Dense │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_12/Dense_2 │ Dense │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_12/Dense_3 │ Dense │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_12/Dense_4 │ Dense │ float32[4,256,1024] │ bfloat16[4,256,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_12/LayerNorm_1 │ LayerNorm │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_12/MlpBlock_0 │ MlpBlock │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_12/MlpBlock_0/Dense_0 │ Dense │ bfloat16[4,256,1024] │ bfloat16[4,256,4096] │ bias: float32[4096] │ +│ │ │ │ │ kernel: float32[1024,4096] │ +│ │ │ │ │ │ +│ │ │ │ │ 4,198,400 (16.8 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_12/MlpBlock_0/Dropout_0 │ Dropout │ bfloat16[4,256,4096] │ bfloat16[4,256,4096] │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_12/MlpBlock_0/Dense_1 │ Dense │ bfloat16[4,256,4096] │ bfloat16[4,256,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[4096,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 4,195,328 (16.8 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_12/MlpBlock_0/Dropout_1 │ Dropout │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_13 │ DiTBlock │ - bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ │ +│ │ │ - float32[1,1024] │ │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_13/Dense_0 │ Dense │ float32[1,1024] │ bfloat16[1,6144] │ bias: float32[6144] │ +│ │ │ │ │ kernel: float32[1024,6144] │ +│ │ │ │ │ │ +│ │ │ │ │ 6,297,600 (25.2 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_13/LayerNorm_0 │ LayerNorm │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_13/Dense_1 │ Dense │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_13/Dense_2 │ Dense │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_13/Dense_3 │ Dense │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_13/Dense_4 │ Dense │ float32[4,256,1024] │ bfloat16[4,256,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_13/LayerNorm_1 │ LayerNorm │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_13/MlpBlock_0 │ MlpBlock │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_13/MlpBlock_0/Dense_0 │ Dense │ bfloat16[4,256,1024] │ bfloat16[4,256,4096] │ bias: float32[4096] │ +│ │ │ │ │ kernel: float32[1024,4096] │ +│ │ │ │ │ │ +│ │ │ │ │ 4,198,400 (16.8 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼──────���────────────────────────┤ +│ DiTBlock_13/MlpBlock_0/Dropout_0 │ Dropout │ bfloat16[4,256,4096] │ bfloat16[4,256,4096] │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_13/MlpBlock_0/Dense_1 │ Dense │ bfloat16[4,256,4096] │ bfloat16[4,256,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[4096,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 4,195,328 (16.8 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_13/MlpBlock_0/Dropout_1 │ Dropout │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_14 │ DiTBlock │ - bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ │ +│ │ │ - float32[1,1024] │ │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_14/Dense_0 │ Dense │ float32[1,1024] │ bfloat16[1,6144] │ bias: float32[6144] │ +│ │ │ │ │ kernel: float32[1024,6144] │ +│ │ │ │ │ │ +│ │ │ │ │ 6,297,600 (25.2 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_14/LayerNorm_0 │ LayerNorm │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_14/Dense_1 │ Dense │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_14/Dense_2 │ Dense │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_14/Dense_3 │ Dense │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_14/Dense_4 │ Dense │ float32[4,256,1024] │ bfloat16[4,256,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_14/LayerNorm_1 │ LayerNorm │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_14/MlpBlock_0 │ MlpBlock │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_14/MlpBlock_0/Dense_0 │ Dense │ bfloat16[4,256,1024] │ bfloat16[4,256,4096] │ bias: float32[4096] │ +│ │ │ │ │ kernel: float32[1024,4096] │ +│ │ │ │ │ │ +│ │ │ │ │ 4,198,400 (16.8 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_14/MlpBlock_0/Dropout_0 │ Dropout │ bfloat16[4,256,4096] │ bfloat16[4,256,4096] │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_14/MlpBlock_0/Dense_1 │ Dense │ bfloat16[4,256,4096] │ bfloat16[4,256,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[4096,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 4,195,328 (16.8 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_14/MlpBlock_0/Dropout_1 │ Dropout │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_15 │ DiTBlock │ - bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ │ +│ │ │ - float32[1,1024] │ │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_15/Dense_0 │ Dense │ float32[1,1024] │ bfloat16[1,6144] │ bias: float32[6144] │ +│ │ │ │ │ kernel: float32[1024,6144] │ +│ │ │ │ │ │ +│ │ │ │ │ 6,297,600 (25.2 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_15/LayerNorm_0 │ LayerNorm │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_15/Dense_1 │ Dense │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_15/Dense_2 │ Dense │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_15/Dense_3 │ Dense │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_15/Dense_4 │ Dense │ float32[4,256,1024] │ bfloat16[4,256,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_15/LayerNorm_1 │ LayerNorm │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼─��──────────────────────┼───────────────────────────────┤ +│ DiTBlock_15/MlpBlock_0 │ MlpBlock │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_15/MlpBlock_0/Dense_0 │ Dense │ bfloat16[4,256,1024] │ bfloat16[4,256,4096] │ bias: float32[4096] │ +│ │ │ │ │ kernel: float32[1024,4096] │ +│ │ │ │ │ │ +│ │ │ │ │ 4,198,400 (16.8 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_15/MlpBlock_0/Dropout_0 │ Dropout │ bfloat16[4,256,4096] │ bfloat16[4,256,4096] │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_15/MlpBlock_0/Dense_1 │ Dense │ bfloat16[4,256,4096] │ bfloat16[4,256,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[4096,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 4,195,328 (16.8 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_15/MlpBlock_0/Dropout_1 │ Dropout │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_16 │ DiTBlock │ - bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ │ +│ │ │ - float32[1,1024] │ │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_16/Dense_0 │ Dense │ float32[1,1024] │ bfloat16[1,6144] │ bias: float32[6144] │ +│ │ │ │ │ kernel: float32[1024,6144] │ +│ │ │ │ │ │ +│ │ │ │ │ 6,297,600 (25.2 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_16/LayerNorm_0 │ LayerNorm │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_16/Dense_1 │ Dense │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_16/Dense_2 │ Dense │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_16/Dense_3 │ Dense │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_16/Dense_4 │ Dense │ float32[4,256,1024] │ bfloat16[4,256,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_16/LayerNorm_1 │ LayerNorm │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_16/MlpBlock_0 │ MlpBlock │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_16/MlpBlock_0/Dense_0 │ Dense │ bfloat16[4,256,1024] │ bfloat16[4,256,4096] │ bias: float32[4096] │ +│ │ │ │ │ kernel: float32[1024,4096] │ +│ │ │ │ │ │ +│ │ │ │ │ 4,198,400 (16.8 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_16/MlpBlock_0/Dropout_0 │ Dropout │ bfloat16[4,256,4096] │ bfloat16[4,256,4096] │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_16/MlpBlock_0/Dense_1 │ Dense │ bfloat16[4,256,4096] │ bfloat16[4,256,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[4096,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 4,195,328 (16.8 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_16/MlpBlock_0/Dropout_1 │ Dropout │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_17 │ DiTBlock │ - bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ │ +│ │ │ - float32[1,1024] │ │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_17/Dense_0 │ Dense │ float32[1,1024] │ bfloat16[1,6144] │ bias: float32[6144] │ +│ │ │ │ │ kernel: float32[1024,6144] │ +│ │ │ │ │ │ +│ │ │ │ │ 6,297,600 (25.2 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_17/LayerNorm_0 │ LayerNorm │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_17/Dense_1 │ Dense │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_17/Dense_2 │ Dense │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_17/Dense_3 │ Dense │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_17/Dense_4 │ Dense │ float32[4,256,1024] │ bfloat16[4,256,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_17/LayerNorm_1 │ LayerNorm │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_17/MlpBlock_0 │ MlpBlock │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_17/MlpBlock_0/Dense_0 │ Dense │ bfloat16[4,256,1024] │ bfloat16[4,256,4096] │ bias: float32[4096] │ +│ │ │ │ │ kernel: float32[1024,4096] │ +│ │ │ │ │ │ +│ │ │ │ │ 4,198,400 (16.8 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_17/MlpBlock_0/Dropout_0 │ Dropout │ bfloat16[4,256,4096] │ bfloat16[4,256,4096] │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_17/MlpBlock_0/Dense_1 │ Dense │ bfloat16[4,256,4096] │ bfloat16[4,256,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[4096,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 4,195,328 (16.8 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼──────��────────────────────────┤ +│ DiTBlock_17/MlpBlock_0/Dropout_1 │ Dropout │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_18 │ DiTBlock │ - bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ │ +│ │ │ - float32[1,1024] │ │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_18/Dense_0 │ Dense │ float32[1,1024] │ bfloat16[1,6144] │ bias: float32[6144] │ +│ │ │ │ │ kernel: float32[1024,6144] │ +│ │ │ │ │ │ +│ │ │ │ │ 6,297,600 (25.2 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_18/LayerNorm_0 │ LayerNorm │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_18/Dense_1 │ Dense │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_18/Dense_2 │ Dense │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼─────────────���─────────────────┤ +│ DiTBlock_18/Dense_3 │ Dense │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_18/Dense_4 │ Dense │ float32[4,256,1024] │ bfloat16[4,256,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_18/LayerNorm_1 │ LayerNorm │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_18/MlpBlock_0 │ MlpBlock │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_18/MlpBlock_0/Dense_0 │ Dense │ bfloat16[4,256,1024] │ bfloat16[4,256,4096] │ bias: float32[4096] │ +│ │ │ │ │ kernel: float32[1024,4096] │ +│ │ │ │ │ │ +│ │ │ │ │ 4,198,400 (16.8 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_18/MlpBlock_0/Dropout_0 │ Dropout │ bfloat16[4,256,4096] │ bfloat16[4,256,4096] │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_18/MlpBlock_0/Dense_1 │ Dense │ bfloat16[4,256,4096] │ bfloat16[4,256,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[4096,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 4,195,328 (16.8 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_18/MlpBlock_0/Dropout_1 │ Dropout │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_19 │ DiTBlock │ - bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ │ +│ │ │ - float32[1,1024] │ │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_19/Dense_0 │ Dense │ float32[1,1024] │ bfloat16[1,6144] │ bias: float32[6144] │ +│ │ │ │ │ kernel: float32[1024,6144] │ +│ │ │ │ │ │ +│ │ │ │ │ 6,297,600 (25.2 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_19/LayerNorm_0 │ LayerNorm │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_19/Dense_1 │ Dense │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_19/Dense_2 │ Dense │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_19/Dense_3 │ Dense │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_19/Dense_4 │ Dense │ float32[4,256,1024] │ bfloat16[4,256,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_19/LayerNorm_1 │ LayerNorm │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_19/MlpBlock_0 │ MlpBlock │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_19/MlpBlock_0/Dense_0 │ Dense │ bfloat16[4,256,1024] │ bfloat16[4,256,4096] │ bias: float32[4096] │ +│ │ │ │ │ kernel: float32[1024,4096] │ +│ │ │ │ │ │ +│ │ │ │ │ 4,198,400 (16.8 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_19/MlpBlock_0/Dropout_0 │ Dropout │ bfloat16[4,256,4096] │ bfloat16[4,256,4096] │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_19/MlpBlock_0/Dense_1 │ Dense │ bfloat16[4,256,4096] │ bfloat16[4,256,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[4096,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 4,195,328 (16.8 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_19/MlpBlock_0/Dropout_1 │ Dropout │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_20 │ DiTBlock │ - bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ │ +│ │ │ - float32[1,1024] │ │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_20/Dense_0 │ Dense │ float32[1,1024] │ bfloat16[1,6144] │ bias: float32[6144] │ +│ │ │ │ │ kernel: float32[1024,6144] │ +│ │ │ │ │ │ +│ │ │ │ │ 6,297,600 (25.2 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_20/LayerNorm_0 │ LayerNorm │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_20/Dense_1 │ Dense │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_20/Dense_2 │ Dense │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_20/Dense_3 │ Dense │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_20/Dense_4 │ Dense │ float32[4,256,1024] │ bfloat16[4,256,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_20/LayerNorm_1 │ LayerNorm │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_20/MlpBlock_0 │ MlpBlock │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_20/MlpBlock_0/Dense_0 │ Dense │ bfloat16[4,256,1024] │ bfloat16[4,256,4096] │ bias: float32[4096] │ +│ │ │ │ │ kernel: float32[1024,4096] │ +│ │ │ │ │ │ +│ │ │ │ │ 4,198,400 (16.8 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_20/MlpBlock_0/Dropout_0 │ Dropout │ bfloat16[4,256,4096] │ bfloat16[4,256,4096] │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_20/MlpBlock_0/Dense_1 │ Dense │ bfloat16[4,256,4096] │ bfloat16[4,256,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[4096,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 4,195,328 (16.8 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_20/MlpBlock_0/Dropout_1 │ Dropout │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_21 │ DiTBlock │ - bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ │ +│ │ │ - float32[1,1024] │ │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_21/Dense_0 │ Dense │ float32[1,1024] │ bfloat16[1,6144] │ bias: float32[6144] │ +│ │ │ │ │ kernel: float32[1024,6144] │ +│ │ │ │ │ │ +│ │ │ │ │ 6,297,600 (25.2 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_21/LayerNorm_0 │ LayerNorm │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_21/Dense_1 │ Dense │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_21/Dense_2 │ Dense │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_21/Dense_3 │ Dense │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_21/Dense_4 │ Dense │ float32[4,256,1024] │ bfloat16[4,256,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_21/LayerNorm_1 │ LayerNorm │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ │ +├──────────────────────────────────┼────────────────���─┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_21/MlpBlock_0 │ MlpBlock │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_21/MlpBlock_0/Dense_0 │ Dense │ bfloat16[4,256,1024] │ bfloat16[4,256,4096] │ bias: float32[4096] │ +│ │ │ │ │ kernel: float32[1024,4096] │ +│ │ │ │ │ │ +│ │ │ │ │ 4,198,400 (16.8 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_21/MlpBlock_0/Dropout_0 │ Dropout │ bfloat16[4,256,4096] │ bfloat16[4,256,4096] │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_21/MlpBlock_0/Dense_1 │ Dense │ bfloat16[4,256,4096] │ bfloat16[4,256,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[4096,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 4,195,328 (16.8 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_21/MlpBlock_0/Dropout_1 │ Dropout │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_22 │ DiTBlock │ - bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ │ +│ │ │ - float32[1,1024] │ │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_22/Dense_0 │ Dense │ float32[1,1024] │ bfloat16[1,6144] │ bias: float32[6144] │ +│ │ │ │ │ kernel: float32[1024,6144] │ +│ │ │ │ │ │ +│ │ │ │ │ 6,297,600 (25.2 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_22/LayerNorm_0 │ LayerNorm │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_22/Dense_1 │ Dense │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_22/Dense_2 │ Dense │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_22/Dense_3 │ Dense │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_22/Dense_4 │ Dense │ float32[4,256,1024] │ bfloat16[4,256,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ ��� │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_22/LayerNorm_1 │ LayerNorm │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_22/MlpBlock_0 │ MlpBlock │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_22/MlpBlock_0/Dense_0 │ Dense │ bfloat16[4,256,1024] │ bfloat16[4,256,4096] │ bias: float32[4096] │ +│ │ │ │ │ kernel: float32[1024,4096] │ +│ │ │ │ │ │ +│ │ │ │ │ 4,198,400 (16.8 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_22/MlpBlock_0/Dropout_0 │ Dropout │ bfloat16[4,256,4096] │ bfloat16[4,256,4096] │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_22/MlpBlock_0/Dense_1 │ Dense │ bfloat16[4,256,4096] │ bfloat16[4,256,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[4096,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 4,195,328 (16.8 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_22/MlpBlock_0/Dropout_1 │ Dropout │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_23 │ DiTBlock │ - bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ │ +│ │ │ - float32[1,1024] │ │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_23/Dense_0 │ Dense │ float32[1,1024] │ bfloat16[1,6144] │ bias: float32[6144] │ +│ │ │ │ │ kernel: float32[1024,6144] │ +│ │ │ │ │ │ +│ │ │ │ │ 6,297,600 (25.2 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_23/LayerNorm_0 │ LayerNorm │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_23/Dense_1 │ Dense │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_23/Dense_2 │ Dense │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_23/Dense_3 │ Dense │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├─────────────────────��────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_23/Dense_4 │ Dense │ float32[4,256,1024] │ bfloat16[4,256,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_23/LayerNorm_1 │ LayerNorm │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_23/MlpBlock_0 │ MlpBlock │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_23/MlpBlock_0/Dense_0 │ Dense │ bfloat16[4,256,1024] │ bfloat16[4,256,4096] │ bias: float32[4096] │ +│ │ │ │ │ kernel: float32[1024,4096] │ +│ │ │ │ │ │ +│ │ │ │ │ 4,198,400 (16.8 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_23/MlpBlock_0/Dropout_0 │ Dropout │ bfloat16[4,256,4096] │ bfloat16[4,256,4096] │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ DiTBlock_23/MlpBlock_0/Dense_1 │ Dense │ bfloat16[4,256,4096] │ bfloat16[4,256,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[4096,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 4,195,328 (16.8 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼──���─────────────────────┼───────────────────────────────┤ +│ DiTBlock_23/MlpBlock_0/Dropout_1 │ Dropout │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ FinalLayer_0 │ FinalLayer │ - bfloat16[4,256,1024] │ bfloat16[4,256,16] │ │ +│ │ │ - float32[1,1024] │ │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ FinalLayer_0/Dense_0 │ Dense │ float32[1,1024] │ bfloat16[1,2048] │ bias: float32[2048] │ +│ │ │ │ │ kernel: float32[1024,2048] │ +│ │ │ │ │ │ +│ │ │ │ │ 2,099,200 (8.4 MB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ FinalLayer_0/LayerNorm_0 │ LayerNorm │ bfloat16[4,256,1024] │ bfloat16[4,256,1024] │ │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ FinalLayer_0/Dense_1 │ Dense │ bfloat16[4,256,1024] │ bfloat16[4,256,16] │ bias: float32[16] │ +│ │ │ │ │ kernel: float32[1024,16] │ +│ │ │ │ │ │ +│ │ │ │ │ 16,400 (65.6 KB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│ Embed_0 │ Embed │ int32[1] │ float32[1,1] │ embedding: float32[256,1] │ +│ │ │ │ │ │ +│ │ │ │ │ 256 (1.0 KB) │ +├──────────────────────────────────┼──────────────────┼────────────────────────┼────────────────────────┼───────────────────────────────┤ +│   │   │   │  Total │ 459,137,296 (1.8 GB)  │ +└──────────────────────────────────┴──────────────────┴────────────────────────┴────────────────────────┴───────────────────────────────┘ +  + Total Parameters: 459,137,296 (1.8 GB)  + + +DiT: Input of shape (4, 32, 32, 4) dtype float32 +DiT: After patch embed, shape is (4, 256, 1024) dtype bfloat16 +DiT: Patch Embed of shape (4, 256, 1024) dtype bfloat16 +DiT: Conditioning of shape (1, 1024) dtype float32 +Loaded checkpoint from 37236 seconds ago. + + parameter shapes: +('DiTBlock_0', 'Dense_0', 'bias'): (1, 6144) +('DiTBlock_0', 'Dense_0', 'kernel'): (1, 1024, 6144) +('DiTBlock_0', 'Dense_1', 'bias'): (1, 1024) +('DiTBlock_0', 'Dense_1', 'kernel'): (1, 1024, 1024) +('DiTBlock_0', 'Dense_2', 'bias'): (1, 1024) +('DiTBlock_0', 'Dense_2', 'kernel'): (1, 1024, 1024) +('DiTBlock_0', 'Dense_3', 'bias'): (1, 1024) +('DiTBlock_0', 'Dense_3', 'kernel'): (1, 1024, 1024) +('DiTBlock_0', 'Dense_4', 'bias'): (1, 1024) +('DiTBlock_0', 'Dense_4', 'kernel'): (1, 1024, 1024) +('DiTBlock_0', 'MlpBlock_0', 'Dense_0', 'bias'): (1, 4096) +('DiTBlock_0', 'MlpBlock_0', 'Dense_0', 'kernel'): (1, 1024, 4096) +('DiTBlock_0', 'MlpBlock_0', 'Dense_1', 'bias'): (1, 1024) +('DiTBlock_0', 'MlpBlock_0', 'Dense_1', 'kernel'): (1, 4096, 1024) +('DiTBlock_1', 'Dense_0', 'bias'): (1, 6144) +('DiTBlock_1', 'Dense_0', 'kernel'): (1, 1024, 6144) +('DiTBlock_1', 'Dense_1', 'bias'): (1, 1024) +('DiTBlock_1', 'Dense_1', 'kernel'): (1, 1024, 1024) +('DiTBlock_1', 'Dense_2', 'bias'): (1, 1024) +('DiTBlock_1', 'Dense_2', 'kernel'): (1, 1024, 1024) +('DiTBlock_1', 'Dense_3', 'bias'): (1, 1024) +('DiTBlock_1', 'Dense_3', 'kernel'): (1, 1024, 1024) +('DiTBlock_1', 'Dense_4', 'bias'): (1, 1024) +('DiTBlock_1', 'Dense_4', 'kernel'): (1, 1024, 1024) +('DiTBlock_1', 'MlpBlock_0', 'Dense_0', 'bias'): (1, 4096) +('DiTBlock_1', 'MlpBlock_0', 'Dense_0', 'kernel'): (1, 1024, 4096) +('DiTBlock_1', 'MlpBlock_0', 'Dense_1', 'bias'): (1, 1024) +('DiTBlock_1', 'MlpBlock_0', 'Dense_1', 'kernel'): (1, 4096, 1024) +('DiTBlock_10', 'Dense_0', 'bias'): (1, 6144) +('DiTBlock_10', 'Dense_0', 'kernel'): (1, 1024, 6144) +('DiTBlock_10', 'Dense_1', 'bias'): (1, 1024) +('DiTBlock_10', 'Dense_1', 'kernel'): (1, 1024, 1024) +('DiTBlock_10', 'Dense_2', 'bias'): (1, 1024) +('DiTBlock_10', 'Dense_2', 'kernel'): (1, 1024, 1024) +('DiTBlock_10', 'Dense_3', 'bias'): (1, 1024) +('DiTBlock_10', 'Dense_3', 'kernel'): (1, 1024, 1024) +('DiTBlock_10', 'Dense_4', 'bias'): (1, 1024) +('DiTBlock_10', 'Dense_4', 'kernel'): (1, 1024, 1024) +('DiTBlock_10', 'MlpBlock_0', 'Dense_0', 'bias'): (1, 4096) +('DiTBlock_10', 'MlpBlock_0', 'Dense_0', 'kernel'): (1, 1024, 4096) +('DiTBlock_10', 'MlpBlock_0', 'Dense_1', 'bias'): (1, 1024) +('DiTBlock_10', 'MlpBlock_0', 'Dense_1', 'kernel'): (1, 4096, 1024) +('DiTBlock_11', 'Dense_0', 'bias'): (1, 6144) +('DiTBlock_11', 'Dense_0', 'kernel'): (1, 1024, 6144) +('DiTBlock_11', 'Dense_1', 'bias'): (1, 1024) +('DiTBlock_11', 'Dense_1', 'kernel'): (1, 1024, 1024) +('DiTBlock_11', 'Dense_2', 'bias'): (1, 1024) +('DiTBlock_11', 'Dense_2', 'kernel'): (1, 1024, 1024) +('DiTBlock_11', 'Dense_3', 'bias'): (1, 1024) +('DiTBlock_11', 'Dense_3', 'kernel'): (1, 1024, 1024) +('DiTBlock_11', 'Dense_4', 'bias'): (1, 1024) +('DiTBlock_11', 'Dense_4', 'kernel'): (1, 1024, 1024) +('DiTBlock_11', 'MlpBlock_0', 'Dense_0', 'bias'): (1, 4096) +('DiTBlock_11', 'MlpBlock_0', 'Dense_0', 'kernel'): (1, 1024, 4096) +('DiTBlock_11', 'MlpBlock_0', 'Dense_1', 'bias'): (1, 1024) +('DiTBlock_11', 'MlpBlock_0', 'Dense_1', 'kernel'): (1, 4096, 1024) +('DiTBlock_12', 'Dense_0', 'bias'): (1, 6144) +('DiTBlock_12', 'Dense_0', 'kernel'): (1, 1024, 6144) +('DiTBlock_12', 'Dense_1', 'bias'): (1, 1024) +('DiTBlock_12', 'Dense_1', 'kernel'): (1, 1024, 1024) +('DiTBlock_12', 'Dense_2', 'bias'): (1, 1024) +('DiTBlock_12', 'Dense_2', 'kernel'): (1, 1024, 1024) +('DiTBlock_12', 'Dense_3', 'bias'): (1, 1024) +('DiTBlock_12', 'Dense_3', 'kernel'): (1, 1024, 1024) +('DiTBlock_12', 'Dense_4', 'bias'): (1, 1024) +('DiTBlock_12', 'Dense_4', 'kernel'): (1, 1024, 1024) +('DiTBlock_12', 'MlpBlock_0', 'Dense_0', 'bias'): (1, 4096) +('DiTBlock_12', 'MlpBlock_0', 'Dense_0', 'kernel'): (1, 1024, 4096) +('DiTBlock_12', 'MlpBlock_0', 'Dense_1', 'bias'): (1, 1024) +('DiTBlock_12', 'MlpBlock_0', 'Dense_1', 'kernel'): (1, 4096, 1024) +('DiTBlock_13', 'Dense_0', 'bias'): (1, 6144) +('DiTBlock_13', 'Dense_0', 'kernel'): (1, 1024, 6144) +('DiTBlock_13', 'Dense_1', 'bias'): (1, 1024) +('DiTBlock_13', 'Dense_1', 'kernel'): (1, 1024, 1024) +('DiTBlock_13', 'Dense_2', 'bias'): (1, 1024) +('DiTBlock_13', 'Dense_2', 'kernel'): (1, 1024, 1024) +('DiTBlock_13', 'Dense_3', 'bias'): (1, 1024) +('DiTBlock_13', 'Dense_3', 'kernel'): (1, 1024, 1024) +('DiTBlock_13', 'Dense_4', 'bias'): (1, 1024) +('DiTBlock_13', 'Dense_4', 'kernel'): (1, 1024, 1024) +('DiTBlock_13', 'MlpBlock_0', 'Dense_0', 'bias'): (1, 4096) +('DiTBlock_13', 'MlpBlock_0', 'Dense_0', 'kernel'): (1, 1024, 4096) +('DiTBlock_13', 'MlpBlock_0', 'Dense_1', 'bias'): (1, 1024) +('DiTBlock_13', 'MlpBlock_0', 'Dense_1', 'kernel'): (1, 4096, 1024) +('DiTBlock_14', 'Dense_0', 'bias'): (1, 6144) +('DiTBlock_14', 'Dense_0', 'kernel'): (1, 1024, 6144) +('DiTBlock_14', 'Dense_1', 'bias'): (1, 1024) +('DiTBlock_14', 'Dense_1', 'kernel'): (1, 1024, 1024) +('DiTBlock_14', 'Dense_2', 'bias'): (1, 1024) +('DiTBlock_14', 'Dense_2', 'kernel'): (1, 1024, 1024) +('DiTBlock_14', 'Dense_3', 'bias'): (1, 1024) +('DiTBlock_14', 'Dense_3', 'kernel'): (1, 1024, 1024) +('DiTBlock_14', 'Dense_4', 'bias'): (1, 1024) +('DiTBlock_14', 'Dense_4', 'kernel'): (1, 1024, 1024) +('DiTBlock_14', 'MlpBlock_0', 'Dense_0', 'bias'): (1, 4096) +('DiTBlock_14', 'MlpBlock_0', 'Dense_0', 'kernel'): (1, 1024, 4096) +('DiTBlock_14', 'MlpBlock_0', 'Dense_1', 'bias'): (1, 1024) +('DiTBlock_14', 'MlpBlock_0', 'Dense_1', 'kernel'): (1, 4096, 1024) +('DiTBlock_15', 'Dense_0', 'bias'): (1, 6144) +('DiTBlock_15', 'Dense_0', 'kernel'): (1, 1024, 6144) +('DiTBlock_15', 'Dense_1', 'bias'): (1, 1024) +('DiTBlock_15', 'Dense_1', 'kernel'): (1, 1024, 1024) +('DiTBlock_15', 'Dense_2', 'bias'): (1, 1024) +('DiTBlock_15', 'Dense_2', 'kernel'): (1, 1024, 1024) +('DiTBlock_15', 'Dense_3', 'bias'): (1, 1024) +('DiTBlock_15', 'Dense_3', 'kernel'): (1, 1024, 1024) +('DiTBlock_15', 'Dense_4', 'bias'): (1, 1024) +('DiTBlock_15', 'Dense_4', 'kernel'): (1, 1024, 1024) +('DiTBlock_15', 'MlpBlock_0', 'Dense_0', 'bias'): (1, 4096) +('DiTBlock_15', 'MlpBlock_0', 'Dense_0', 'kernel'): (1, 1024, 4096) +('DiTBlock_15', 'MlpBlock_0', 'Dense_1', 'bias'): (1, 1024) +('DiTBlock_15', 'MlpBlock_0', 'Dense_1', 'kernel'): (1, 4096, 1024) +('DiTBlock_16', 'Dense_0', 'bias'): (1, 6144) +('DiTBlock_16', 'Dense_0', 'kernel'): (1, 1024, 6144) +('DiTBlock_16', 'Dense_1', 'bias'): (1, 1024) +('DiTBlock_16', 'Dense_1', 'kernel'): (1, 1024, 1024) +('DiTBlock_16', 'Dense_2', 'bias'): (1, 1024) +('DiTBlock_16', 'Dense_2', 'kernel'): (1, 1024, 1024) +('DiTBlock_16', 'Dense_3', 'bias'): (1, 1024) +('DiTBlock_16', 'Dense_3', 'kernel'): (1, 1024, 1024) +('DiTBlock_16', 'Dense_4', 'bias'): (1, 1024) +('DiTBlock_16', 'Dense_4', 'kernel'): (1, 1024, 1024) +('DiTBlock_16', 'MlpBlock_0', 'Dense_0', 'bias'): (1, 4096) +('DiTBlock_16', 'MlpBlock_0', 'Dense_0', 'kernel'): (1, 1024, 4096) +('DiTBlock_16', 'MlpBlock_0', 'Dense_1', 'bias'): (1, 1024) +('DiTBlock_16', 'MlpBlock_0', 'Dense_1', 'kernel'): (1, 4096, 1024) +('DiTBlock_17', 'Dense_0', 'bias'): (1, 6144) +('DiTBlock_17', 'Dense_0', 'kernel'): (1, 1024, 6144) +('DiTBlock_17', 'Dense_1', 'bias'): (1, 1024) +('DiTBlock_17', 'Dense_1', 'kernel'): (1, 1024, 1024) +('DiTBlock_17', 'Dense_2', 'bias'): (1, 1024) +('DiTBlock_17', 'Dense_2', 'kernel'): (1, 1024, 1024) +('DiTBlock_17', 'Dense_3', 'bias'): (1, 1024) +('DiTBlock_17', 'Dense_3', 'kernel'): (1, 1024, 1024) +('DiTBlock_17', 'Dense_4', 'bias'): (1, 1024) +('DiTBlock_17', 'Dense_4', 'kernel'): (1, 1024, 1024) +('DiTBlock_17', 'MlpBlock_0', 'Dense_0', 'bias'): (1, 4096) +('DiTBlock_17', 'MlpBlock_0', 'Dense_0', 'kernel'): (1, 1024, 4096) +('DiTBlock_17', 'MlpBlock_0', 'Dense_1', 'bias'): (1, 1024) +('DiTBlock_17', 'MlpBlock_0', 'Dense_1', 'kernel'): (1, 4096, 1024) +('DiTBlock_18', 'Dense_0', 'bias'): (1, 6144) +('DiTBlock_18', 'Dense_0', 'kernel'): (1, 1024, 6144) +('DiTBlock_18', 'Dense_1', 'bias'): (1, 1024) +('DiTBlock_18', 'Dense_1', 'kernel'): (1, 1024, 1024) +('DiTBlock_18', 'Dense_2', 'bias'): (1, 1024) +('DiTBlock_18', 'Dense_2', 'kernel'): (1, 1024, 1024) +('DiTBlock_18', 'Dense_3', 'bias'): (1, 1024) +('DiTBlock_18', 'Dense_3', 'kernel'): (1, 1024, 1024) +('DiTBlock_18', 'Dense_4', 'bias'): (1, 1024) +('DiTBlock_18', 'Dense_4', 'kernel'): (1, 1024, 1024) +('DiTBlock_18', 'MlpBlock_0', 'Dense_0', 'bias'): (1, 4096) +('DiTBlock_18', 'MlpBlock_0', 'Dense_0', 'kernel'): (1, 1024, 4096) +('DiTBlock_18', 'MlpBlock_0', 'Dense_1', 'bias'): (1, 1024) +('DiTBlock_18', 'MlpBlock_0', 'Dense_1', 'kernel'): (1, 4096, 1024) +('DiTBlock_19', 'Dense_0', 'bias'): (1, 6144) +('DiTBlock_19', 'Dense_0', 'kernel'): (1, 1024, 6144) +('DiTBlock_19', 'Dense_1', 'bias'): (1, 1024) +('DiTBlock_19', 'Dense_1', 'kernel'): (1, 1024, 1024) +('DiTBlock_19', 'Dense_2', 'bias'): (1, 1024) +('DiTBlock_19', 'Dense_2', 'kernel'): (1, 1024, 1024) +('DiTBlock_19', 'Dense_3', 'bias'): (1, 1024) +('DiTBlock_19', 'Dense_3', 'kernel'): (1, 1024, 1024) +('DiTBlock_19', 'Dense_4', 'bias'): (1, 1024) +('DiTBlock_19', 'Dense_4', 'kernel'): (1, 1024, 1024) +('DiTBlock_19', 'MlpBlock_0', 'Dense_0', 'bias'): (1, 4096) +('DiTBlock_19', 'MlpBlock_0', 'Dense_0', 'kernel'): (1, 1024, 4096) +('DiTBlock_19', 'MlpBlock_0', 'Dense_1', 'bias'): (1, 1024) +('DiTBlock_19', 'MlpBlock_0', 'Dense_1', 'kernel'): (1, 4096, 1024) +('DiTBlock_2', 'Dense_0', 'bias'): (1, 6144) +('DiTBlock_2', 'Dense_0', 'kernel'): (1, 1024, 6144) +('DiTBlock_2', 'Dense_1', 'bias'): (1, 1024) +('DiTBlock_2', 'Dense_1', 'kernel'): (1, 1024, 1024) +('DiTBlock_2', 'Dense_2', 'bias'): (1, 1024) +('DiTBlock_2', 'Dense_2', 'kernel'): (1, 1024, 1024) +('DiTBlock_2', 'Dense_3', 'bias'): (1, 1024) +('DiTBlock_2', 'Dense_3', 'kernel'): (1, 1024, 1024) +('DiTBlock_2', 'Dense_4', 'bias'): (1, 1024) +('DiTBlock_2', 'Dense_4', 'kernel'): (1, 1024, 1024) +('DiTBlock_2', 'MlpBlock_0', 'Dense_0', 'bias'): (1, 4096) +('DiTBlock_2', 'MlpBlock_0', 'Dense_0', 'kernel'): (1, 1024, 4096) +('DiTBlock_2', 'MlpBlock_0', 'Dense_1', 'bias'): (1, 1024) +('DiTBlock_2', 'MlpBlock_0', 'Dense_1', 'kernel'): (1, 4096, 1024) +('DiTBlock_20', 'Dense_0', 'bias'): (1, 6144) +('DiTBlock_20', 'Dense_0', 'kernel'): (1, 1024, 6144) +('DiTBlock_20', 'Dense_1', 'bias'): (1, 1024) +('DiTBlock_20', 'Dense_1', 'kernel'): (1, 1024, 1024) +('DiTBlock_20', 'Dense_2', 'bias'): (1, 1024) +('DiTBlock_20', 'Dense_2', 'kernel'): (1, 1024, 1024) +('DiTBlock_20', 'Dense_3', 'bias'): (1, 1024) +('DiTBlock_20', 'Dense_3', 'kernel'): (1, 1024, 1024) +('DiTBlock_20', 'Dense_4', 'bias'): (1, 1024) +('DiTBlock_20', 'Dense_4', 'kernel'): (1, 1024, 1024) +('DiTBlock_20', 'MlpBlock_0', 'Dense_0', 'bias'): (1, 4096) +('DiTBlock_20', 'MlpBlock_0', 'Dense_0', 'kernel'): (1, 1024, 4096) +('DiTBlock_20', 'MlpBlock_0', 'Dense_1', 'bias'): (1, 1024) +('DiTBlock_20', 'MlpBlock_0', 'Dense_1', 'kernel'): (1, 4096, 1024) +('DiTBlock_21', 'Dense_0', 'bias'): (1, 6144) +('DiTBlock_21', 'Dense_0', 'kernel'): (1, 1024, 6144) +('DiTBlock_21', 'Dense_1', 'bias'): (1, 1024) +('DiTBlock_21', 'Dense_1', 'kernel'): (1, 1024, 1024) +('DiTBlock_21', 'Dense_2', 'bias'): (1, 1024) +('DiTBlock_21', 'Dense_2', 'kernel'): (1, 1024, 1024) +('DiTBlock_21', 'Dense_3', 'bias'): (1, 1024) +('DiTBlock_21', 'Dense_3', 'kernel'): (1, 1024, 1024) +('DiTBlock_21', 'Dense_4', 'bias'): (1, 1024) +('DiTBlock_21', 'Dense_4', 'kernel'): (1, 1024, 1024) +('DiTBlock_21', 'MlpBlock_0', 'Dense_0', 'bias'): (1, 4096) +('DiTBlock_21', 'MlpBlock_0', 'Dense_0', 'kernel'): (1, 1024, 4096) +('DiTBlock_21', 'MlpBlock_0', 'Dense_1', 'bias'): (1, 1024) +('DiTBlock_21', 'MlpBlock_0', 'Dense_1', 'kernel'): (1, 4096, 1024) +('DiTBlock_22', 'Dense_0', 'bias'): (1, 6144) +('DiTBlock_22', 'Dense_0', 'kernel'): (1, 1024, 6144) +('DiTBlock_22', 'Dense_1', 'bias'): (1, 1024) +('DiTBlock_22', 'Dense_1', 'kernel'): (1, 1024, 1024) +('DiTBlock_22', 'Dense_2', 'bias'): (1, 1024) +('DiTBlock_22', 'Dense_2', 'kernel'): (1, 1024, 1024) +('DiTBlock_22', 'Dense_3', 'bias'): (1, 1024) +('DiTBlock_22', 'Dense_3', 'kernel'): (1, 1024, 1024) +('DiTBlock_22', 'Dense_4', 'bias'): (1, 1024) +('DiTBlock_22', 'Dense_4', 'kernel'): (1, 1024, 1024) +('DiTBlock_22', 'MlpBlock_0', 'Dense_0', 'bias'): (1, 4096) +('DiTBlock_22', 'MlpBlock_0', 'Dense_0', 'kernel'): (1, 1024, 4096) +('DiTBlock_22', 'MlpBlock_0', 'Dense_1', 'bias'): (1, 1024) +('DiTBlock_22', 'MlpBlock_0', 'Dense_1', 'kernel'): (1, 4096, 1024) +('DiTBlock_23', 'Dense_0', 'bias'): (1, 6144) +('DiTBlock_23', 'Dense_0', 'kernel'): (1, 1024, 6144) +('DiTBlock_23', 'Dense_1', 'bias'): (1, 1024) +('DiTBlock_23', 'Dense_1', 'kernel'): (1, 1024, 1024) +('DiTBlock_23', 'Dense_2', 'bias'): (1, 1024) +('DiTBlock_23', 'Dense_2', 'kernel'): (1, 1024, 1024) +('DiTBlock_23', 'Dense_3', 'bias'): (1, 1024) +('DiTBlock_23', 'Dense_3', 'kernel'): (1, 1024, 1024) +('DiTBlock_23', 'Dense_4', 'bias'): (1, 1024) +('DiTBlock_23', 'Dense_4', 'kernel'): (1, 1024, 1024) +('DiTBlock_23', 'MlpBlock_0', 'Dense_0', 'bias'): (1, 4096) +('DiTBlock_23', 'MlpBlock_0', 'Dense_0', 'kernel'): (1, 1024, 4096) +('DiTBlock_23', 'MlpBlock_0', 'Dense_1', 'bias'): (1, 1024) +('DiTBlock_23', 'MlpBlock_0', 'Dense_1', 'kernel'): (1, 4096, 1024) +('DiTBlock_3', 'Dense_0', 'bias'): (1, 6144) +('DiTBlock_3', 'Dense_0', 'kernel'): (1, 1024, 6144) +('DiTBlock_3', 'Dense_1', 'bias'): (1, 1024) +('DiTBlock_3', 'Dense_1', 'kernel'): (1, 1024, 1024) +('DiTBlock_3', 'Dense_2', 'bias'): (1, 1024) +('DiTBlock_3', 'Dense_2', 'kernel'): (1, 1024, 1024) +('DiTBlock_3', 'Dense_3', 'bias'): (1, 1024) +('DiTBlock_3', 'Dense_3', 'kernel'): (1, 1024, 1024) +('DiTBlock_3', 'Dense_4', 'bias'): (1, 1024) +('DiTBlock_3', 'Dense_4', 'kernel'): (1, 1024, 1024) +('DiTBlock_3', 'MlpBlock_0', 'Dense_0', 'bias'): (1, 4096) +('DiTBlock_3', 'MlpBlock_0', 'Dense_0', 'kernel'): (1, 1024, 4096) +('DiTBlock_3', 'MlpBlock_0', 'Dense_1', 'bias'): (1, 1024) +('DiTBlock_3', 'MlpBlock_0', 'Dense_1', 'kernel'): (1, 4096, 1024) +('DiTBlock_4', 'Dense_0', 'bias'): (1, 6144) +('DiTBlock_4', 'Dense_0', 'kernel'): (1, 1024, 6144) +('DiTBlock_4', 'Dense_1', 'bias'): (1, 1024) +('DiTBlock_4', 'Dense_1', 'kernel'): (1, 1024, 1024) +('DiTBlock_4', 'Dense_2', 'bias'): (1, 1024) +('DiTBlock_4', 'Dense_2', 'kernel'): (1, 1024, 1024) +('DiTBlock_4', 'Dense_3', 'bias'): (1, 1024) +('DiTBlock_4', 'Dense_3', 'kernel'): (1, 1024, 1024) +('DiTBlock_4', 'Dense_4', 'bias'): (1, 1024) +('DiTBlock_4', 'Dense_4', 'kernel'): (1, 1024, 1024) +('DiTBlock_4', 'MlpBlock_0', 'Dense_0', 'bias'): (1, 4096) +('DiTBlock_4', 'MlpBlock_0', 'Dense_0', 'kernel'): (1, 1024, 4096) +('DiTBlock_4', 'MlpBlock_0', 'Dense_1', 'bias'): (1, 1024) +('DiTBlock_4', 'MlpBlock_0', 'Dense_1', 'kernel'): (1, 4096, 1024) +('DiTBlock_5', 'Dense_0', 'bias'): (1, 6144) +('DiTBlock_5', 'Dense_0', 'kernel'): (1, 1024, 6144) +('DiTBlock_5', 'Dense_1', 'bias'): (1, 1024) +('DiTBlock_5', 'Dense_1', 'kernel'): (1, 1024, 1024) +('DiTBlock_5', 'Dense_2', 'bias'): (1, 1024) +('DiTBlock_5', 'Dense_2', 'kernel'): (1, 1024, 1024) +('DiTBlock_5', 'Dense_3', 'bias'): (1, 1024) +('DiTBlock_5', 'Dense_3', 'kernel'): (1, 1024, 1024) +('DiTBlock_5', 'Dense_4', 'bias'): (1, 1024) +('DiTBlock_5', 'Dense_4', 'kernel'): (1, 1024, 1024) +('DiTBlock_5', 'MlpBlock_0', 'Dense_0', 'bias'): (1, 4096) +('DiTBlock_5', 'MlpBlock_0', 'Dense_0', 'kernel'): (1, 1024, 4096) +('DiTBlock_5', 'MlpBlock_0', 'Dense_1', 'bias'): (1, 1024) +('DiTBlock_5', 'MlpBlock_0', 'Dense_1', 'kernel'): (1, 4096, 1024) +('DiTBlock_6', 'Dense_0', 'bias'): (1, 6144) +('DiTBlock_6', 'Dense_0', 'kernel'): (1, 1024, 6144) +('DiTBlock_6', 'Dense_1', 'bias'): (1, 1024) +('DiTBlock_6', 'Dense_1', 'kernel'): (1, 1024, 1024) +('DiTBlock_6', 'Dense_2', 'bias'): (1, 1024) +('DiTBlock_6', 'Dense_2', 'kernel'): (1, 1024, 1024) +('DiTBlock_6', 'Dense_3', 'bias'): (1, 1024) +('DiTBlock_6', 'Dense_3', 'kernel'): (1, 1024, 1024) +('DiTBlock_6', 'Dense_4', 'bias'): (1, 1024) +('DiTBlock_6', 'Dense_4', 'kernel'): (1, 1024, 1024) +('DiTBlock_6', 'MlpBlock_0', 'Dense_0', 'bias'): (1, 4096) +('DiTBlock_6', 'MlpBlock_0', 'Dense_0', 'kernel'): (1, 1024, 4096) +('DiTBlock_6', 'MlpBlock_0', 'Dense_1', 'bias'): (1, 1024) +('DiTBlock_6', 'MlpBlock_0', 'Dense_1', 'kernel'): (1, 4096, 1024) +('DiTBlock_7', 'Dense_0', 'bias'): (1, 6144) +('DiTBlock_7', 'Dense_0', 'kernel'): (1, 1024, 6144) +('DiTBlock_7', 'Dense_1', 'bias'): (1, 1024) +('DiTBlock_7', 'Dense_1', 'kernel'): (1, 1024, 1024) +('DiTBlock_7', 'Dense_2', 'bias'): (1, 1024) +('DiTBlock_7', 'Dense_2', 'kernel'): (1, 1024, 1024) +('DiTBlock_7', 'Dense_3', 'bias'): (1, 1024) +('DiTBlock_7', 'Dense_3', 'kernel'): (1, 1024, 1024) +('DiTBlock_7', 'Dense_4', 'bias'): (1, 1024) +('DiTBlock_7', 'Dense_4', 'kernel'): (1, 1024, 1024) +('DiTBlock_7', 'MlpBlock_0', 'Dense_0', 'bias'): (1, 4096) +('DiTBlock_7', 'MlpBlock_0', 'Dense_0', 'kernel'): (1, 1024, 4096) +('DiTBlock_7', 'MlpBlock_0', 'Dense_1', 'bias'): (1, 1024) +('DiTBlock_7', 'MlpBlock_0', 'Dense_1', 'kernel'): (1, 4096, 1024) +('DiTBlock_8', 'Dense_0', 'bias'): (1, 6144) +('DiTBlock_8', 'Dense_0', 'kernel'): (1, 1024, 6144) +('DiTBlock_8', 'Dense_1', 'bias'): (1, 1024) +('DiTBlock_8', 'Dense_1', 'kernel'): (1, 1024, 1024) +('DiTBlock_8', 'Dense_2', 'bias'): (1, 1024) +('DiTBlock_8', 'Dense_2', 'kernel'): (1, 1024, 1024) +('DiTBlock_8', 'Dense_3', 'bias'): (1, 1024) +('DiTBlock_8', 'Dense_3', 'kernel'): (1, 1024, 1024) +('DiTBlock_8', 'Dense_4', 'bias'): (1, 1024) +('DiTBlock_8', 'Dense_4', 'kernel'): (1, 1024, 1024) +('DiTBlock_8', 'MlpBlock_0', 'Dense_0', 'bias'): (1, 4096) +('DiTBlock_8', 'MlpBlock_0', 'Dense_0', 'kernel'): (1, 1024, 4096) +('DiTBlock_8', 'MlpBlock_0', 'Dense_1', 'bias'): (1, 1024) +('DiTBlock_8', 'MlpBlock_0', 'Dense_1', 'kernel'): (1, 4096, 1024) +('DiTBlock_9', 'Dense_0', 'bias'): (1, 6144) +('DiTBlock_9', 'Dense_0', 'kernel'): (1, 1024, 6144) +('DiTBlock_9', 'Dense_1', 'bias'): (1, 1024) +('DiTBlock_9', 'Dense_1', 'kernel'): (1, 1024, 1024) +('DiTBlock_9', 'Dense_2', 'bias'): (1, 1024) +('DiTBlock_9', 'Dense_2', 'kernel'): (1, 1024, 1024) +('DiTBlock_9', 'Dense_3', 'bias'): (1, 1024) +('DiTBlock_9', 'Dense_3', 'kernel'): (1, 1024, 1024) +('DiTBlock_9', 'Dense_4', 'bias'): (1, 1024) +('DiTBlock_9', 'Dense_4', 'kernel'): (1, 1024, 1024) +('DiTBlock_9', 'MlpBlock_0', 'Dense_0', 'bias'): (1, 4096) +('DiTBlock_9', 'MlpBlock_0', 'Dense_0', 'kernel'): (1, 1024, 4096) +('DiTBlock_9', 'MlpBlock_0', 'Dense_1', 'bias'): (1, 1024) +('DiTBlock_9', 'MlpBlock_0', 'Dense_1', 'kernel'): (1, 4096, 1024) +('Embed_0', 'embedding'): (1, 256, 1) +('FinalLayer_0', 'Dense_0', 'bias'): (1, 2048) +('FinalLayer_0', 'Dense_0', 'kernel'): (1, 1024, 2048) +('FinalLayer_0', 'Dense_1', 'bias'): (1, 16) +('FinalLayer_0', 'Dense_1', 'kernel'): (1, 1024, 16) +('LabelEmbedder_0', 'Embed_0', 'embedding'): (1, 1001, 1024) +('PatchEmbed_0', 'Conv_0', 'bias'): (1, 1024) +('PatchEmbed_0', 'Conv_0', 'kernel'): (1, 2, 2, 4, 1024) +('TimestepEmbedder_0', 'Dense_0', 'bias'): (1, 1024) +('TimestepEmbedder_0', 'Dense_0', 'kernel'): (1, 256, 1024) +('TimestepEmbedder_0', 'Dense_1', 'bias'): (1, 1024) +('TimestepEmbedder_0', 'Dense_1', 'kernel'): (1, 1024, 1024) +('TimestepEmbedder_1', 'Dense_0', 'bias'): (1, 1024) +('TimestepEmbedder_1', 'Dense_0', 'kernel'): (1, 256, 1024) +('TimestepEmbedder_1', 'Dense_1', 'bias'): (1, 1024) +('TimestepEmbedder_1', 'Dense_1', 'kernel'): (1, 1024, 1024) +NamedSharding(mesh=Mesh('devices': 4, axis_types=(Auto,)), spec=PartitionSpec(), memory_kind=device) +did stuff +┌────────────────────────────────────────────────┐ +│ │ +│ │ +│ │ +│ │ +│ 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, 4) +After quant (128, 32, 32, 4) +Calc FID for CFG 1.0 and denoise_timesteps 128 +DiT: Input of shape (512, 32, 32, 4) dtype float32 +DiT: After patch embed, shape is (512, 256, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 256, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) 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 18.711158752441406 +(512, 256, 256, 3) +Calc FID for CFG 1.0 and denoise_timesteps 64 +DiT: Input of shape (512, 32, 32, 4) dtype float32 +DiT: After patch embed, shape is (512, 256, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 256, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 19.22484588623047 +(512, 256, 256, 3) +Calc FID for CFG 1.0 and denoise_timesteps 32 +DiT: Input of shape (512, 32, 32, 4) dtype float32 +DiT: After patch embed, shape is (512, 256, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 256, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 20.787307739257812 +(512, 256, 256, 3) +Calc FID for CFG 1.0 and denoise_timesteps 16 +DiT: Input of shape (512, 32, 32, 4) dtype float32 +DiT: After patch embed, shape is (512, 256, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 256, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 25.715614318847656 +(512, 256, 256, 3) +Calc FID for CFG 1.0 and denoise_timesteps 8 +DiT: Input of shape (512, 32, 32, 4) dtype float32 +DiT: After patch embed, shape is (512, 256, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 256, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 41.448814392089844 +(512, 256, 256, 3) +Calc FID for CFG 1.0 and denoise_timesteps 4 +DiT: Input of shape (512, 32, 32, 4) dtype float32 +DiT: After patch embed, shape is (512, 256, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 256, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 98.53426361083984 +(512, 256, 256, 3) +Calc FID for CFG 1.0 and denoise_timesteps 2 +DiT: Input of shape (512, 32, 32, 4) dtype float32 +DiT: After patch embed, shape is (512, 256, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 256, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 301.80596923828125 +(512, 256, 256, 3) +Calc FID for CFG 1.0 and denoise_timesteps 1 +DiT: Input of shape (512, 32, 32, 4) dtype float32 +DiT: After patch embed, shape is (512, 256, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 256, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 297.51544189453125 +(512, 256, 256, 3) +Calc FID for CFG 1.25 and denoise_timesteps 128 +DiT: Input of shape (512, 32, 32, 4) dtype float32 +DiT: After patch embed, shape is (512, 256, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 256, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 9.431501388549805 +(512, 256, 256, 3) +Calc FID for CFG 1.25 and denoise_timesteps 64 +DiT: Input of shape (512, 32, 32, 4) dtype float32 +DiT: After patch embed, shape is (512, 256, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 256, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 9.721562385559082 +(512, 256, 256, 3) +Calc FID for CFG 1.25 and denoise_timesteps 32 +DiT: Input of shape (512, 32, 32, 4) dtype float32 +DiT: After patch embed, shape is (512, 256, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 256, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 10.677177429199219 +(512, 256, 256, 3) +Calc FID for CFG 1.25 and denoise_timesteps 16 +DiT: Input of shape (512, 32, 32, 4) dtype float32 +DiT: After patch embed, shape is (512, 256, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 256, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 13.967844009399414 +(512, 256, 256, 3) +Calc FID for CFG 1.25 and denoise_timesteps 8 +DiT: Input of shape (512, 32, 32, 4) dtype float32 +DiT: After patch embed, shape is (512, 256, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 256, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 25.67169189453125 +(512, 256, 256, 3) +Calc FID for CFG 1.25 and denoise_timesteps 4 +DiT: Input of shape (512, 32, 32, 4) dtype float32 +DiT: After patch embed, shape is (512, 256, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 256, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 73.56874084472656 +(512, 256, 256, 3) +Calc FID for CFG 1.25 and denoise_timesteps 2 +DiT: Input of shape (512, 32, 32, 4) dtype float32 +DiT: After patch embed, shape is (512, 256, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 256, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 282.5548095703125 +(512, 256, 256, 3) +Calc FID for CFG 1.25 and denoise_timesteps 1 +DiT: Input of shape (512, 32, 32, 4) dtype float32 +DiT: After patch embed, shape is (512, 256, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 256, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 287.9568786621094 +(512, 256, 256, 3) +Calc FID for CFG 1.5 and denoise_timesteps 128 +DiT: Input of shape (512, 32, 32, 4) dtype float32 +DiT: After patch embed, shape is (512, 256, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 256, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 6.008155822753906 +(512, 256, 256, 3) +Calc FID for CFG 1.5 and denoise_timesteps 64 +DiT: Input of shape (512, 32, 32, 4) dtype float32 +DiT: After patch embed, shape is (512, 256, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 256, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 6.1163249015808105 +(512, 256, 256, 3) +Calc FID for CFG 1.5 and denoise_timesteps 32 +DiT: Input of shape (512, 32, 32, 4) dtype float32 +DiT: After patch embed, shape is (512, 256, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 256, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 6.643142223358154 +(512, 256, 256, 3) +Calc FID for CFG 1.5 and denoise_timesteps 16 +DiT: Input of shape (512, 32, 32, 4) dtype float32 +DiT: After patch embed, shape is (512, 256, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 256, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 8.617753982543945 +(512, 256, 256, 3) +Calc FID for CFG 1.5 and denoise_timesteps 8 +DiT: Input of shape (512, 32, 32, 4) dtype float32 +DiT: After patch embed, shape is (512, 256, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 256, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 16.720661163330078 +(512, 256, 256, 3) +Calc FID for CFG 1.5 and denoise_timesteps 4 +DiT: Input of shape (512, 32, 32, 4) dtype float32 +DiT: After patch embed, shape is (512, 256, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 256, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 55.4674072265625 +(512, 256, 256, 3) +Calc FID for CFG 1.5 and denoise_timesteps 2 +DiT: Input of shape (512, 32, 32, 4) dtype float32 +DiT: After patch embed, shape is (512, 256, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 256, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 263.9443664550781 +(512, 256, 256, 3) +Calc FID for CFG 1.5 and denoise_timesteps 1 +DiT: Input of shape (512, 32, 32, 4) dtype float32 +DiT: After patch embed, shape is (512, 256, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 256, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 278.1593017578125 +(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, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 256, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 5.549402236938477 +(512, 256, 256, 3) +Calc FID for CFG 1.75 and denoise_timesteps 64 +DiT: Input of shape (512, 32, 32, 4) dtype float32 +DiT: After patch embed, shape is (512, 256, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 256, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 5.594891548156738 +(512, 256, 256, 3) +Calc FID for CFG 1.75 and denoise_timesteps 32 +DiT: Input of shape (512, 32, 32, 4) dtype float32 +DiT: After patch embed, shape is (512, 256, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 256, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 5.83141565322876 +(512, 256, 256, 3) +Calc FID for CFG 1.75 and denoise_timesteps 16 +DiT: Input of shape (512, 32, 32, 4) dtype float32 +DiT: After patch embed, shape is (512, 256, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 256, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 6.923707008361816 +(512, 256, 256, 3) +Calc FID for CFG 1.75 and denoise_timesteps 8 +DiT: Input of shape (512, 32, 32, 4) dtype float32 +DiT: After patch embed, shape is (512, 256, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 256, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 12.29236125946045 +(512, 256, 256, 3) +Calc FID for CFG 1.75 and denoise_timesteps 4 +DiT: Input of shape (512, 32, 32, 4) dtype float32 +DiT: After patch embed, shape is (512, 256, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 256, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 42.79435729980469 +(512, 256, 256, 3) +Calc FID for CFG 1.75 and denoise_timesteps 2 +DiT: Input of shape (512, 32, 32, 4) dtype float32 +DiT: After patch embed, shape is (512, 256, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 256, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 247.306396484375 +(512, 256, 256, 3) +Calc FID for CFG 1.75 and denoise_timesteps 1 +DiT: Input of shape (512, 32, 32, 4) dtype float32 +DiT: After patch embed, shape is (512, 256, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 256, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 269.83770751953125 +(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, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 256, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 6.50559139251709 +(512, 256, 256, 3) +Calc FID for CFG 2.0 and denoise_timesteps 64 +DiT: Input of shape (512, 32, 32, 4) dtype float32 +DiT: After patch embed, shape is (512, 256, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 256, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 6.476692199707031 +(512, 256, 256, 3) +Calc FID for CFG 2.0 and denoise_timesteps 32 +DiT: Input of shape (512, 32, 32, 4) dtype float32 +DiT: After patch embed, shape is (512, 256, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 256, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 6.533879280090332 +(512, 256, 256, 3) +Calc FID for CFG 2.0 and denoise_timesteps 16 +DiT: Input of shape (512, 32, 32, 4) dtype float32 +DiT: After patch embed, shape is (512, 256, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 256, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 7.073963642120361 +(512, 256, 256, 3) +Calc FID for CFG 2.0 and denoise_timesteps 8 +DiT: Input of shape (512, 32, 32, 4) dtype float32 +DiT: After patch embed, shape is (512, 256, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 256, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 10.455595016479492 +(512, 256, 256, 3) +Calc FID for CFG 2.0 and denoise_timesteps 4 +DiT: Input of shape (512, 32, 32, 4) dtype float32 +DiT: After patch embed, shape is (512, 256, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 256, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 34.20777893066406 +(512, 256, 256, 3) +Calc FID for CFG 2.0 and denoise_timesteps 2 +DiT: Input of shape (512, 32, 32, 4) dtype float32 +DiT: After patch embed, shape is (512, 256, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 256, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 233.12899780273438 +(512, 256, 256, 3) +Calc FID for CFG 2.0 and denoise_timesteps 1 +DiT: Input of shape (512, 32, 32, 4) dtype float32 +DiT: After patch embed, shape is (512, 256, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 256, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 262.98980712890625 +(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, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 256, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 8.025604248046875 +(512, 256, 256, 3) +Calc FID for CFG 2.25 and denoise_timesteps 64 +DiT: Input of shape (512, 32, 32, 4) dtype float32 +DiT: After patch embed, shape is (512, 256, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 256, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 7.9627461433410645 +(512, 256, 256, 3) +Calc FID for CFG 2.25 and denoise_timesteps 32 +DiT: Input of shape (512, 32, 32, 4) dtype float32 +DiT: After patch embed, shape is (512, 256, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 256, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 7.963242530822754 +(512, 256, 256, 3) +Calc FID for CFG 2.25 and denoise_timesteps 16 +DiT: Input of shape (512, 32, 32, 4) dtype float32 +DiT: After patch embed, shape is (512, 256, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 256, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 8.124591827392578 +(512, 256, 256, 3) +Calc FID for CFG 2.25 and denoise_timesteps 8 +DiT: Input of shape (512, 32, 32, 4) dtype float32 +DiT: After patch embed, shape is (512, 256, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 256, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 10.11518669128418 +(512, 256, 256, 3) +Calc FID for CFG 2.25 and denoise_timesteps 4 +DiT: Input of shape (512, 32, 32, 4) dtype float32 +DiT: After patch embed, shape is (512, 256, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 256, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 28.450166702270508 +(512, 256, 256, 3) +Calc FID for CFG 2.25 and denoise_timesteps 2 +DiT: Input of shape (512, 32, 32, 4) dtype float32 +DiT: After patch embed, shape is (512, 256, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 256, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 221.21920776367188 +(512, 256, 256, 3) +Calc FID for CFG 2.25 and denoise_timesteps 1 +DiT: Input of shape (512, 32, 32, 4) dtype float32 +DiT: After patch embed, shape is (512, 256, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 256, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 256.34344482421875 +(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, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 256, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 9.693035125732422 +(512, 256, 256, 3) +Calc FID for CFG 2.5 and denoise_timesteps 64 +DiT: Input of shape (512, 32, 32, 4) dtype float32 +DiT: After patch embed, shape is (512, 256, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 256, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 9.645331382751465 +(512, 256, 256, 3) +Calc FID for CFG 2.5 and denoise_timesteps 32 +DiT: Input of shape (512, 32, 32, 4) dtype float32 +DiT: After patch embed, shape is (512, 256, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 256, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 9.554620742797852 +(512, 256, 256, 3) +Calc FID for CFG 2.5 and denoise_timesteps 16 +DiT: Input of shape (512, 32, 32, 4) dtype float32 +DiT: After patch embed, shape is (512, 256, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 256, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 9.504511833190918 +(512, 256, 256, 3) +Calc FID for CFG 2.5 and denoise_timesteps 8 +DiT: Input of shape (512, 32, 32, 4) dtype float32 +DiT: After patch embed, shape is (512, 256, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 256, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 10.5525484085083 +(512, 256, 256, 3) +Calc FID for CFG 2.5 and denoise_timesteps 4 +DiT: Input of shape (512, 32, 32, 4) dtype float32 +DiT: After patch embed, shape is (512, 256, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 256, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 24.626758575439453 +(512, 256, 256, 3) +Calc FID for CFG 2.5 and denoise_timesteps 2 +DiT: Input of shape (512, 32, 32, 4) dtype float32 +DiT: After patch embed, shape is (512, 256, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 256, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 211.2801513671875 +(512, 256, 256, 3) +Calc FID for CFG 2.5 and denoise_timesteps 1 +DiT: Input of shape (512, 32, 32, 4) dtype float32 +DiT: After patch embed, shape is (512, 256, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 256, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 251.06690979003906 +(512, 256, 256, 3) +Calc FID for CFG 2.75 and denoise_timesteps 128 +DiT: Input of shape (512, 32, 32, 4) dtype float32 +DiT: After patch embed, shape is (512, 256, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 256, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 11.351409912109375 +(512, 256, 256, 3) +Calc FID for CFG 2.75 and denoise_timesteps 64 +DiT: Input of shape (512, 32, 32, 4) dtype float32 +DiT: After patch embed, shape is (512, 256, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 256, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 11.261353492736816 +(512, 256, 256, 3) +Calc FID for CFG 2.75 and denoise_timesteps 32 +DiT: Input of shape (512, 32, 32, 4) dtype float32 +DiT: After patch embed, shape is (512, 256, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 256, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 11.144073486328125 +(512, 256, 256, 3) +Calc FID for CFG 2.75 and denoise_timesteps 16 +DiT: Input of shape (512, 32, 32, 4) dtype float32 +DiT: After patch embed, shape is (512, 256, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 256, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 10.952659606933594 +(512, 256, 256, 3) +Calc FID for CFG 2.75 and denoise_timesteps 8 +DiT: Input of shape (512, 32, 32, 4) dtype float32 +DiT: After patch embed, shape is (512, 256, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 256, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 11.270659446716309 +(512, 256, 256, 3) +Calc FID for CFG 2.75 and denoise_timesteps 4 +DiT: Input of shape (512, 32, 32, 4) dtype float32 +DiT: After patch embed, shape is (512, 256, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 256, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 22.154510498046875 +(512, 256, 256, 3) +Calc FID for CFG 2.75 and denoise_timesteps 2 +DiT: Input of shape (512, 32, 32, 4) dtype float32 +DiT: After patch embed, shape is (512, 256, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 256, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 202.94566345214844 +(512, 256, 256, 3) +Calc FID for CFG 2.75 and denoise_timesteps 1 +DiT: Input of shape (512, 32, 32, 4) dtype float32 +DiT: After patch embed, shape is (512, 256, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 256, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 247.4989471435547 +(512, 256, 256, 3) +Calc FID for CFG 3.0 and denoise_timesteps 128 +DiT: Input of shape (512, 32, 32, 4) dtype float32 +DiT: After patch embed, shape is (512, 256, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 256, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 12.887807846069336 +(512, 256, 256, 3) +Calc FID for CFG 3.0 and denoise_timesteps 64 +DiT: Input of shape (512, 32, 32, 4) dtype float32 +DiT: After patch embed, shape is (512, 256, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 256, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 12.796567916870117 +(512, 256, 256, 3) +Calc FID for CFG 3.0 and denoise_timesteps 32 +DiT: Input of shape (512, 32, 32, 4) dtype float32 +DiT: After patch embed, shape is (512, 256, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 256, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 12.658254623413086 +(512, 256, 256, 3) +Calc FID for CFG 3.0 and denoise_timesteps 16 +DiT: Input of shape (512, 32, 32, 4) dtype float32 +DiT: After patch embed, shape is (512, 256, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 256, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 12.399557113647461 +(512, 256, 256, 3) +Calc FID for CFG 3.0 and denoise_timesteps 8 +DiT: Input of shape (512, 32, 32, 4) dtype float32 +DiT: After patch embed, shape is (512, 256, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 256, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 12.165553092956543 +(512, 256, 256, 3) +Calc FID for CFG 3.0 and denoise_timesteps 4 +DiT: Input of shape (512, 32, 32, 4) dtype float32 +DiT: After patch embed, shape is (512, 256, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 256, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 20.568153381347656 +(512, 256, 256, 3) +Calc FID for CFG 3.0 and denoise_timesteps 2 +DiT: Input of shape (512, 32, 32, 4) dtype float32 +DiT: After patch embed, shape is (512, 256, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 256, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 196.02975463867188 +(512, 256, 256, 3) +Calc FID for CFG 3.0 and denoise_timesteps 1 +DiT: Input of shape (512, 32, 32, 4) dtype float32 +DiT: After patch embed, shape is (512, 256, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 256, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 244.86386108398438