diff --git "a/f16_c16_large/logEma.txt" "b/f16_c16_large/logEma.txt" new file mode 100644--- /dev/null +++ "b/f16_c16_large/logEma.txt" @@ -0,0 +1,2263 @@ +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, 128) +doing downsample +Encoder layer (1, 32, 32, 256) +doing downsample +Encoder layer (1, 16, 16, 256) +Encoder layer (1, 16, 16, 512) +Encoder layer final (1, 16, 16, 512) +Encoder layer final (1, 16, 16, 512) +Final embeddings are size (1, 16, 16, 16) +After quant (1, 16, 16, 16) +encode finished +Decoder incoming shape (1, 16, 16, 16) +Decoder input (1, 16, 16, 512) +Mid Block Decoder layer (1, 16, 16, 512) +Mid Block Decoder layer (1, 16, 16, 512) +Decoder layer (1, 32, 32, 512) +Decoder layer (1, 64, 64, 256) +Decoder layer (1, 128, 128, 256) +Decoder layer (1, 256, 256, 128) +Decoder layer (1, 256, 256, 128) +Total num of VQVAE parameters: 53024403 +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 1157331 seconds ago. +Loaded model with step 989001 +┌──────────────────────────────────────────────────────────────────────────────┐ +│ 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, 128) +doing downsample +Encoder layer (1, 32, 32, 256) +doing downsample +Encoder layer (1, 16, 16, 256) +Encoder layer (1, 16, 16, 512) +Encoder layer final (1, 16, 16, 512) +Encoder layer final (1, 16, 16, 512) +Final embeddings are size (1, 16, 16, 16) +After quant (1, 16, 16, 16) +output example shape (4, 1, 16, 16, 16) +Test data shape (4, 256, 256, 3) +x shape (4, 1, 256, 256, 3) +encoded shape (4, 1, 16, 16, 16) +z_vectors shape (1, 16, 16, 16) +Decoder incoming shape (1, 16, 16, 16) +Decoder input (1, 16, 16, 512) +Mid Block Decoder layer (1, 16, 16, 512) +Mid Block Decoder layer (1, 16, 16, 512) +Decoder layer (1, 32, 32, 512) +Decoder layer (1, 64, 64, 256) +Decoder layer (1, 128, 128, 256) +Decoder layer (1, 256, 256, 128) +Decoder layer (1, 256, 256, 128) +image shape (4, 1, 256, 256, 3) +decoded img shape (256, 256, 3) +obs shape (4, 16, 16, 16) +DiT: Input of shape (4, 16, 16, 16) dtype float32 +DiT: After patch embed, shape is (4, 64, 1024) dtype bfloat16 +DiT: Patch Embed of shape (4, 64, 1024) dtype bfloat16 +DiT: Conditioning of shape (1, 1024) dtype float32 + + DiT Summary  +┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓ +┃ path  ┃ module  ┃ inputs  ┃ outputs  ┃ params  ┃ +┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┩ +│ │ DiT │ - float32[4,16,16,16] │ bfloat16[4,16,16,16] │ │ +│ │ │ - float32[1] │ │ │ +│ │ │ - float32[1] │ │ │ +│ │ │ - int32[1] │ │ │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ PatchEmbed_0 │ PatchEmbed │ float32[4,16,16,16] │ bfloat16[4,64,1024] │ │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ PatchEmbed_0/Conv_0 │ Conv │ float32[4,16,16,16] │ bfloat16[4,8,8,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[2,2,16,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 66,560 (266.2 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,64,1024] │ bfloat16[4,64,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,64,1024] │ bfloat16[4,64,1024] │ │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_0/Dense_1 │ Dense │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_0/Dense_2 │ Dense │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_0/Dense_3 │ Dense │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_0/Dense_4 │ Dense │ float32[4,64,1024] │ bfloat16[4,64,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_0/LayerNorm_1 │ LayerNorm │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_0/MlpBlock_0 │ MlpBlock │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_0/MlpBlock_0/Dense_0 │ Dense │ bfloat16[4,64,1024] │ bfloat16[4,64,4096] │ bias: float32[4096] │ +│ │ │ │ │ kernel: float32[1024,4096] │ +│ │ │ │ │ │ +│ │ │ │ │ 4,198,400 (16.8 MB) │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_0/MlpBlock_0/Dropout_0 │ Dropout │ bfloat16[4,64,4096] │ bfloat16[4,64,4096] │ │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_0/MlpBlock_0/Dense_1 │ Dense │ bfloat16[4,64,4096] │ bfloat16[4,64,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[4096,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 4,195,328 (16.8 MB) │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_0/MlpBlock_0/Dropout_1 │ Dropout │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_1 │ DiTBlock │ - bfloat16[4,64,1024] │ bfloat16[4,64,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,64,1024] │ bfloat16[4,64,1024] │ │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_1/Dense_1 │ Dense │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_1/Dense_2 │ Dense │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├───────────────────────────────��──┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_1/Dense_3 │ Dense │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_1/Dense_4 │ Dense │ float32[4,64,1024] │ bfloat16[4,64,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_1/LayerNorm_1 │ LayerNorm │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_1/MlpBlock_0 │ MlpBlock │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_1/MlpBlock_0/Dense_0 │ Dense │ bfloat16[4,64,1024] │ bfloat16[4,64,4096] │ bias: float32[4096] │ +│ │ │ │ │ kernel: float32[1024,4096] │ +│ │ │ │ │ │ +│ │ │ │ │ 4,198,400 (16.8 MB) │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_1/MlpBlock_0/Dropout_0 │ Dropout │ bfloat16[4,64,4096] │ bfloat16[4,64,4096] │ │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────���───────┤ +│ DiTBlock_1/MlpBlock_0/Dense_1 │ Dense │ bfloat16[4,64,4096] │ bfloat16[4,64,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[4096,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 4,195,328 (16.8 MB) │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_1/MlpBlock_0/Dropout_1 │ Dropout │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_2 │ DiTBlock │ - bfloat16[4,64,1024] │ bfloat16[4,64,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,64,1024] │ bfloat16[4,64,1024] │ │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_2/Dense_1 │ Dense │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_2/Dense_2 │ Dense │ bfloat16[4,64,1024] �� bfloat16[4,64,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_2/Dense_3 │ Dense │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_2/Dense_4 │ Dense │ float32[4,64,1024] │ bfloat16[4,64,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_2/LayerNorm_1 │ LayerNorm │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_2/MlpBlock_0 │ MlpBlock │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_2/MlpBlock_0/Dense_0 │ Dense │ bfloat16[4,64,1024] │ bfloat16[4,64,4096] │ bias: float32[4096] │ +│ │ │ │ │ kernel: float32[1024,4096] │ +│ │ │ │ │ │ +│ │ │ │ │ 4,198,400 (16.8 MB) │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼────────────────────���─┼───────────────────────────────┤ +│ DiTBlock_2/MlpBlock_0/Dropout_0 │ Dropout │ bfloat16[4,64,4096] │ bfloat16[4,64,4096] │ │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_2/MlpBlock_0/Dense_1 │ Dense │ bfloat16[4,64,4096] │ bfloat16[4,64,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[4096,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 4,195,328 (16.8 MB) │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_2/MlpBlock_0/Dropout_1 │ Dropout │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_3 │ DiTBlock │ - bfloat16[4,64,1024] │ bfloat16[4,64,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,64,1024] │ bfloat16[4,64,1024] │ │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_3/Dense_1 │ Dense │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_3/Dense_2 │ Dense │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_3/Dense_3 │ Dense │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_3/Dense_4 │ Dense │ float32[4,64,1024] │ bfloat16[4,64,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_3/LayerNorm_1 │ LayerNorm │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_3/MlpBlock_0 │ MlpBlock │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_3/MlpBlock_0/Dense_0 │ Dense │ bfloat16[4,64,1024] │ bfloat16[4,64,4096] │ bias: float32[4096] │ +│ │ │ │ │ kernel: float32[1024,4096] │ +│ │ │ │ │ │ +│ │ │ │ │ 4,198,400 (16.8 MB) │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_3/MlpBlock_0/Dropout_0 │ Dropout │ bfloat16[4,64,4096] │ bfloat16[4,64,4096] │ │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_3/MlpBlock_0/Dense_1 │ Dense │ bfloat16[4,64,4096] │ bfloat16[4,64,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[4096,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 4,195,328 (16.8 MB) │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_3/MlpBlock_0/Dropout_1 │ Dropout │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_4 │ DiTBlock │ - bfloat16[4,64,1024] │ bfloat16[4,64,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,64,1024] │ bfloat16[4,64,1024] │ │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_4/Dense_1 │ Dense │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_4/Dense_2 │ Dense │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_4/Dense_3 │ Dense │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_4/Dense_4 │ Dense │ float32[4,64,1024] │ bfloat16[4,64,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_4/LayerNorm_1 │ LayerNorm │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_4/MlpBlock_0 │ MlpBlock │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ │ +├────────────────���─────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_4/MlpBlock_0/Dense_0 │ Dense │ bfloat16[4,64,1024] │ bfloat16[4,64,4096] │ bias: float32[4096] │ +│ │ │ │ │ kernel: float32[1024,4096] │ +│ │ │ │ │ │ +│ │ │ │ │ 4,198,400 (16.8 MB) │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_4/MlpBlock_0/Dropout_0 │ Dropout │ bfloat16[4,64,4096] │ bfloat16[4,64,4096] │ │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_4/MlpBlock_0/Dense_1 │ Dense │ bfloat16[4,64,4096] │ bfloat16[4,64,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[4096,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 4,195,328 (16.8 MB) │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_4/MlpBlock_0/Dropout_1 │ Dropout │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_5 │ DiTBlock │ - bfloat16[4,64,1024] │ bfloat16[4,64,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,64,1024] │ bfloat16[4,64,1024] │ │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_5/Dense_1 │ Dense │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_5/Dense_2 │ Dense │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_5/Dense_3 │ Dense │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_5/Dense_4 │ Dense │ float32[4,64,1024] │ bfloat16[4,64,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_5/LayerNorm_1 │ LayerNorm │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_5/MlpBlock_0 │ MlpBlock │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_5/MlpBlock_0/Dense_0 │ Dense │ bfloat16[4,64,1024] │ bfloat16[4,64,4096] │ bias: float32[4096] │ +│ │ │ │ │ kernel: float32[1024,4096] │ +│ │ │ │ │ │ +│ │ │ │ │ 4,198,400 (16.8 MB) │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_5/MlpBlock_0/Dropout_0 │ Dropout │ bfloat16[4,64,4096] │ bfloat16[4,64,4096] │ │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_5/MlpBlock_0/Dense_1 │ Dense │ bfloat16[4,64,4096] │ bfloat16[4,64,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[4096,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 4,195,328 (16.8 MB) │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_5/MlpBlock_0/Dropout_1 │ Dropout │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_6 │ DiTBlock │ - bfloat16[4,64,1024] │ bfloat16[4,64,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,64,1024] │ bfloat16[4,64,1024] │ │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_6/Dense_1 │ Dense │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_6/Dense_2 │ Dense │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_6/Dense_3 │ Dense │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_6/Dense_4 │ Dense │ float32[4,64,1024] │ bfloat16[4,64,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_6/LayerNorm_1 │ LayerNorm │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_6/MlpBlock_0 │ MlpBlock │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_6/MlpBlock_0/Dense_0 │ Dense │ bfloat16[4,64,1024] │ bfloat16[4,64,4096] │ bias: float32[4096] │ +│ │ │ │ │ kernel: float32[1024,4096] │ +│ │ │ │ │ │ +│ │ │ │ │ 4,198,400 (16.8 MB) │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_6/MlpBlock_0/Dropout_0 │ Dropout │ bfloat16[4,64,4096] │ bfloat16[4,64,4096] │ │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_6/MlpBlock_0/Dense_1 │ Dense │ bfloat16[4,64,4096] │ bfloat16[4,64,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[4096,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 4,195,328 (16.8 MB) │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_6/MlpBlock_0/Dropout_1 │ Dropout │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_7 │ DiTBlock │ - bfloat16[4,64,1024] │ bfloat16[4,64,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,64,1024] │ bfloat16[4,64,1024] │ │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_7/Dense_1 │ Dense │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_7/Dense_2 │ Dense │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_7/Dense_3 │ Dense │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼───────────────────��──┼───────────────────────────────┤ +│ DiTBlock_7/Dense_4 │ Dense │ float32[4,64,1024] │ bfloat16[4,64,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_7/LayerNorm_1 │ LayerNorm │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_7/MlpBlock_0 │ MlpBlock │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_7/MlpBlock_0/Dense_0 │ Dense │ bfloat16[4,64,1024] │ bfloat16[4,64,4096] │ bias: float32[4096] │ +│ │ │ │ │ kernel: float32[1024,4096] │ +│ │ │ │ │ │ +│ │ │ │ │ 4,198,400 (16.8 MB) │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_7/MlpBlock_0/Dropout_0 │ Dropout │ bfloat16[4,64,4096] │ bfloat16[4,64,4096] │ │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_7/MlpBlock_0/Dense_1 │ Dense │ bfloat16[4,64,4096] │ bfloat16[4,64,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[4096,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 4,195,328 (16.8 MB) │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_7/MlpBlock_0/Dropout_1 │ Dropout │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_8 │ DiTBlock │ - bfloat16[4,64,1024] │ bfloat16[4,64,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,64,1024] │ bfloat16[4,64,1024] │ │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_8/Dense_1 │ Dense │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_8/Dense_2 │ Dense │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_8/Dense_3 │ Dense │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_8/Dense_4 │ Dense │ float32[4,64,1024] │ bfloat16[4,64,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_8/LayerNorm_1 │ LayerNorm │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_8/MlpBlock_0 │ MlpBlock │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_8/MlpBlock_0/Dense_0 │ Dense │ bfloat16[4,64,1024] │ bfloat16[4,64,4096] │ bias: float32[4096] │ +│ │ │ │ │ kernel: float32[1024,4096] │ +│ │ │ │ │ │ +│ │ │ │ │ 4,198,400 (16.8 MB) │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_8/MlpBlock_0/Dropout_0 │ Dropout │ bfloat16[4,64,4096] │ bfloat16[4,64,4096] │ │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_8/MlpBlock_0/Dense_1 │ Dense │ bfloat16[4,64,4096] │ bfloat16[4,64,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[4096,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 4,195,328 (16.8 MB) │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_8/MlpBlock_0/Dropout_1 │ Dropout │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_9 │ DiTBlock │ - bfloat16[4,64,1024] │ bfloat16[4,64,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,64,1024] │ bfloat16[4,64,1024] │ │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_9/Dense_1 │ Dense │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_9/Dense_2 │ Dense │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├───────────────��──────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_9/Dense_3 │ Dense │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_9/Dense_4 │ Dense │ float32[4,64,1024] │ bfloat16[4,64,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_9/LayerNorm_1 │ LayerNorm │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_9/MlpBlock_0 │ MlpBlock │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_9/MlpBlock_0/Dense_0 │ Dense │ bfloat16[4,64,1024] │ bfloat16[4,64,4096] │ bias: float32[4096] │ +│ │ │ │ │ kernel: float32[1024,4096] │ +│ │ │ │ │ │ +│ │ │ │ │ 4,198,400 (16.8 MB) │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_9/MlpBlock_0/Dropout_0 │ Dropout │ bfloat16[4,64,4096] │ bfloat16[4,64,4096] │ │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────���───────────────────────┤ +│ DiTBlock_9/MlpBlock_0/Dense_1 │ Dense │ bfloat16[4,64,4096] │ bfloat16[4,64,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[4096,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 4,195,328 (16.8 MB) │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_9/MlpBlock_0/Dropout_1 │ Dropout │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_10 │ DiTBlock │ - bfloat16[4,64,1024] │ bfloat16[4,64,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,64,1024] │ bfloat16[4,64,1024] │ │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_10/Dense_1 │ Dense │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_10/Dense_2 │ Dense │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_10/Dense_3 │ Dense │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_10/Dense_4 │ Dense │ float32[4,64,1024] │ bfloat16[4,64,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_10/LayerNorm_1 │ LayerNorm │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_10/MlpBlock_0 │ MlpBlock │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_10/MlpBlock_0/Dense_0 │ Dense │ bfloat16[4,64,1024] │ bfloat16[4,64,4096] │ bias: float32[4096] │ +│ │ │ │ │ kernel: float32[1024,4096] │ +│ │ │ │ │ │ +│ │ │ │ │ 4,198,400 (16.8 MB) │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼────���─────────────────┼───────────────────────────────┤ +│ DiTBlock_10/MlpBlock_0/Dropout_0 │ Dropout │ bfloat16[4,64,4096] │ bfloat16[4,64,4096] │ │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_10/MlpBlock_0/Dense_1 │ Dense │ bfloat16[4,64,4096] │ bfloat16[4,64,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[4096,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 4,195,328 (16.8 MB) │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_10/MlpBlock_0/Dropout_1 │ Dropout │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_11 │ DiTBlock │ - bfloat16[4,64,1024] │ bfloat16[4,64,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,64,1024] │ bfloat16[4,64,1024] │ │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_11/Dense_1 │ Dense │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_11/Dense_2 │ Dense │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_11/Dense_3 │ Dense │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_11/Dense_4 │ Dense │ float32[4,64,1024] │ bfloat16[4,64,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_11/LayerNorm_1 │ LayerNorm │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_11/MlpBlock_0 │ MlpBlock │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_11/MlpBlock_0/Dense_0 │ Dense │ bfloat16[4,64,1024] │ bfloat16[4,64,4096] │ bias: float32[4096] │ +│ │ │ │ │ kernel: float32[1024,4096] │ +│ │ │ │ │ │ +│ │ │ │ │ 4,198,400 (16.8 MB) │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_11/MlpBlock_0/Dropout_0 │ Dropout │ bfloat16[4,64,4096] │ bfloat16[4,64,4096] │ │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_11/MlpBlock_0/Dense_1 │ Dense │ bfloat16[4,64,4096] │ bfloat16[4,64,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[4096,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 4,195,328 (16.8 MB) │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_11/MlpBlock_0/Dropout_1 │ Dropout │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_12 │ DiTBlock │ - bfloat16[4,64,1024] │ bfloat16[4,64,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,64,1024] │ bfloat16[4,64,1024] │ │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_12/Dense_1 │ Dense │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_12/Dense_2 │ Dense │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_12/Dense_3 │ Dense │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_12/Dense_4 │ Dense │ float32[4,64,1024] │ bfloat16[4,64,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_12/LayerNorm_1 │ LayerNorm │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_12/MlpBlock_0 │ MlpBlock │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ │ +├���─────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_12/MlpBlock_0/Dense_0 │ Dense │ bfloat16[4,64,1024] │ bfloat16[4,64,4096] │ bias: float32[4096] │ +│ │ │ │ │ kernel: float32[1024,4096] │ +│ │ │ │ │ │ +│ │ │ │ │ 4,198,400 (16.8 MB) │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_12/MlpBlock_0/Dropout_0 │ Dropout │ bfloat16[4,64,4096] │ bfloat16[4,64,4096] │ │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_12/MlpBlock_0/Dense_1 │ Dense │ bfloat16[4,64,4096] │ bfloat16[4,64,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[4096,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 4,195,328 (16.8 MB) │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_12/MlpBlock_0/Dropout_1 │ Dropout │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_13 │ DiTBlock │ - bfloat16[4,64,1024] │ bfloat16[4,64,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,64,1024] │ bfloat16[4,64,1024] │ │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_13/Dense_1 │ Dense │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_13/Dense_2 │ Dense │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_13/Dense_3 │ Dense │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_13/Dense_4 │ Dense │ float32[4,64,1024] │ bfloat16[4,64,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_13/LayerNorm_1 │ LayerNorm │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_13/MlpBlock_0 │ MlpBlock │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_13/MlpBlock_0/Dense_0 │ Dense │ bfloat16[4,64,1024] │ bfloat16[4,64,4096] │ bias: float32[4096] │ +│ │ │ │ │ kernel: float32[1024,4096] │ +│ │ │ │ │ │ +│ │ │ │ │ 4,198,400 (16.8 MB) │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_13/MlpBlock_0/Dropout_0 │ Dropout │ bfloat16[4,64,4096] │ bfloat16[4,64,4096] │ │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_13/MlpBlock_0/Dense_1 │ Dense │ bfloat16[4,64,4096] │ bfloat16[4,64,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[4096,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 4,195,328 (16.8 MB) │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_13/MlpBlock_0/Dropout_1 │ Dropout │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_14 │ DiTBlock │ - bfloat16[4,64,1024] │ bfloat16[4,64,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,64,1024] │ bfloat16[4,64,1024] │ │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_14/Dense_1 │ Dense │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_14/Dense_2 │ Dense │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_14/Dense_3 │ Dense │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_14/Dense_4 │ Dense │ float32[4,64,1024] │ bfloat16[4,64,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_14/LayerNorm_1 │ LayerNorm │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_14/MlpBlock_0 │ MlpBlock │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_14/MlpBlock_0/Dense_0 │ Dense │ bfloat16[4,64,1024] │ bfloat16[4,64,4096] │ bias: float32[4096] │ +│ │ │ │ │ kernel: float32[1024,4096] │ +│ │ │ │ │ │ +│ │ │ │ │ 4,198,400 (16.8 MB) │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_14/MlpBlock_0/Dropout_0 │ Dropout │ bfloat16[4,64,4096] │ bfloat16[4,64,4096] │ │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_14/MlpBlock_0/Dense_1 │ Dense │ bfloat16[4,64,4096] │ bfloat16[4,64,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[4096,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 4,195,328 (16.8 MB) │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_14/MlpBlock_0/Dropout_1 │ Dropout │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_15 │ DiTBlock │ - bfloat16[4,64,1024] │ bfloat16[4,64,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,64,1024] │ bfloat16[4,64,1024] │ │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_15/Dense_1 │ Dense │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_15/Dense_2 │ Dense │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_15/Dense_3 │ Dense │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼───��──────────────────┼───────────────────────────────┤ +│ DiTBlock_15/Dense_4 │ Dense │ float32[4,64,1024] │ bfloat16[4,64,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_15/LayerNorm_1 │ LayerNorm │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_15/MlpBlock_0 │ MlpBlock │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_15/MlpBlock_0/Dense_0 │ Dense │ bfloat16[4,64,1024] │ bfloat16[4,64,4096] │ bias: float32[4096] │ +│ │ │ │ │ kernel: float32[1024,4096] │ +│ │ │ │ │ │ +│ │ │ │ │ 4,198,400 (16.8 MB) │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_15/MlpBlock_0/Dropout_0 │ Dropout │ bfloat16[4,64,4096] │ bfloat16[4,64,4096] │ │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_15/MlpBlock_0/Dense_1 │ Dense │ bfloat16[4,64,4096] │ bfloat16[4,64,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[4096,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 4,195,328 (16.8 MB) │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_15/MlpBlock_0/Dropout_1 │ Dropout │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_16 │ DiTBlock │ - bfloat16[4,64,1024] │ bfloat16[4,64,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,64,1024] │ bfloat16[4,64,1024] │ │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_16/Dense_1 │ Dense │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_16/Dense_2 │ Dense │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_16/Dense_3 │ Dense │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_16/Dense_4 │ Dense │ float32[4,64,1024] │ bfloat16[4,64,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_16/LayerNorm_1 │ LayerNorm │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_16/MlpBlock_0 │ MlpBlock │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_16/MlpBlock_0/Dense_0 │ Dense │ bfloat16[4,64,1024] │ bfloat16[4,64,4096] │ bias: float32[4096] │ +│ │ │ │ │ kernel: float32[1024,4096] │ +│ │ │ │ │ │ +│ │ │ │ │ 4,198,400 (16.8 MB) │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_16/MlpBlock_0/Dropout_0 │ Dropout │ bfloat16[4,64,4096] │ bfloat16[4,64,4096] │ │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_16/MlpBlock_0/Dense_1 │ Dense │ bfloat16[4,64,4096] │ bfloat16[4,64,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[4096,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 4,195,328 (16.8 MB) │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_16/MlpBlock_0/Dropout_1 │ Dropout │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_17 │ DiTBlock │ - bfloat16[4,64,1024] │ bfloat16[4,64,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,64,1024] │ bfloat16[4,64,1024] │ │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_17/Dense_1 │ Dense │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_17/Dense_2 │ Dense │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +��──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_17/Dense_3 │ Dense │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_17/Dense_4 │ Dense │ float32[4,64,1024] │ bfloat16[4,64,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_17/LayerNorm_1 │ LayerNorm │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_17/MlpBlock_0 │ MlpBlock │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_17/MlpBlock_0/Dense_0 │ Dense │ bfloat16[4,64,1024] │ bfloat16[4,64,4096] │ bias: float32[4096] │ +│ │ │ │ │ kernel: float32[1024,4096] │ +│ │ │ │ │ │ +│ │ │ │ │ 4,198,400 (16.8 MB) │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_17/MlpBlock_0/Dropout_0 │ Dropout │ bfloat16[4,64,4096] │ bfloat16[4,64,4096] │ │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────���───────┼───────────────────────────────┤ +│ DiTBlock_17/MlpBlock_0/Dense_1 │ Dense │ bfloat16[4,64,4096] │ bfloat16[4,64,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[4096,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 4,195,328 (16.8 MB) │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_17/MlpBlock_0/Dropout_1 │ Dropout │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_18 │ DiTBlock │ - bfloat16[4,64,1024] │ bfloat16[4,64,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,64,1024] │ bfloat16[4,64,1024] │ │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_18/Dense_1 │ Dense │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_18/Dense_2 │ Dense │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_18/Dense_3 │ Dense │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_18/Dense_4 │ Dense │ float32[4,64,1024] │ bfloat16[4,64,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_18/LayerNorm_1 │ LayerNorm │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_18/MlpBlock_0 │ MlpBlock │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_18/MlpBlock_0/Dense_0 │ Dense │ bfloat16[4,64,1024] │ bfloat16[4,64,4096] │ bias: float32[4096] │ +│ │ │ │ │ kernel: float32[1024,4096] │ +│ │ │ │ │ │ +│ │ │ │ │ 4,198,400 (16.8 MB) │ +├──────────────────────────────────┼──────────────────┼────────────���──────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_18/MlpBlock_0/Dropout_0 │ Dropout │ bfloat16[4,64,4096] │ bfloat16[4,64,4096] │ │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_18/MlpBlock_0/Dense_1 │ Dense │ bfloat16[4,64,4096] │ bfloat16[4,64,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[4096,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 4,195,328 (16.8 MB) │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_18/MlpBlock_0/Dropout_1 │ Dropout │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_19 │ DiTBlock │ - bfloat16[4,64,1024] │ bfloat16[4,64,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,64,1024] │ bfloat16[4,64,1024] │ │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_19/Dense_1 │ Dense │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_19/Dense_2 │ Dense │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_19/Dense_3 │ Dense │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_19/Dense_4 │ Dense │ float32[4,64,1024] │ bfloat16[4,64,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_19/LayerNorm_1 │ LayerNorm │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_19/MlpBlock_0 │ MlpBlock │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_19/MlpBlock_0/Dense_0 │ Dense │ bfloat16[4,64,1024] │ bfloat16[4,64,4096] │ bias: float32[4096] │ +│ │ │ │ │ kernel: float32[1024,4096] │ +│ │ │ │ │ │ +│ │ │ │ │ 4,198,400 (16.8 MB) │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_19/MlpBlock_0/Dropout_0 │ Dropout │ bfloat16[4,64,4096] │ bfloat16[4,64,4096] │ │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_19/MlpBlock_0/Dense_1 │ Dense │ bfloat16[4,64,4096] │ bfloat16[4,64,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[4096,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 4,195,328 (16.8 MB) │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_19/MlpBlock_0/Dropout_1 │ Dropout │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_20 │ DiTBlock │ - bfloat16[4,64,1024] │ bfloat16[4,64,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,64,1024] │ bfloat16[4,64,1024] │ │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_20/Dense_1 │ Dense │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_20/Dense_2 │ Dense │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_20/Dense_3 │ Dense │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_20/Dense_4 │ Dense │ float32[4,64,1024] │ bfloat16[4,64,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_20/LayerNorm_1 │ LayerNorm │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_20/MlpBlock_0 │ MlpBlock │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_20/MlpBlock_0/Dense_0 │ Dense │ bfloat16[4,64,1024] │ bfloat16[4,64,4096] │ bias: float32[4096] │ +│ │ │ │ │ kernel: float32[1024,4096] │ +│ │ │ │ │ │ +│ │ │ │ │ 4,198,400 (16.8 MB) │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_20/MlpBlock_0/Dropout_0 │ Dropout │ bfloat16[4,64,4096] │ bfloat16[4,64,4096] │ │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_20/MlpBlock_0/Dense_1 │ Dense │ bfloat16[4,64,4096] │ bfloat16[4,64,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[4096,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 4,195,328 (16.8 MB) │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_20/MlpBlock_0/Dropout_1 │ Dropout │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_21 │ DiTBlock │ - bfloat16[4,64,1024] │ bfloat16[4,64,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,64,1024] │ bfloat16[4,64,1024] │ │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_21/Dense_1 │ Dense │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_21/Dense_2 │ Dense │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_21/Dense_3 │ Dense │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_21/Dense_4 │ Dense │ float32[4,64,1024] │ bfloat16[4,64,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_21/LayerNorm_1 │ LayerNorm │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_21/MlpBlock_0 │ MlpBlock │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_21/MlpBlock_0/Dense_0 │ Dense │ bfloat16[4,64,1024] │ bfloat16[4,64,4096] │ bias: float32[4096] │ +│ │ │ │ │ kernel: float32[1024,4096] │ +│ │ │ │ │ │ +│ │ │ │ │ 4,198,400 (16.8 MB) │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_21/MlpBlock_0/Dropout_0 │ Dropout │ bfloat16[4,64,4096] │ bfloat16[4,64,4096] │ │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_21/MlpBlock_0/Dense_1 │ Dense │ bfloat16[4,64,4096] │ bfloat16[4,64,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[4096,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 4,195,328 (16.8 MB) │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_21/MlpBlock_0/Dropout_1 │ Dropout │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_22 │ DiTBlock │ - bfloat16[4,64,1024] │ bfloat16[4,64,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,64,1024] │ bfloat16[4,64,1024] │ │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_22/Dense_1 │ Dense │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_22/Dense_2 │ Dense │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_22/Dense_3 │ Dense │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_22/Dense_4 │ Dense │ float32[4,64,1024] │ bfloat16[4,64,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_22/LayerNorm_1 │ LayerNorm │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_22/MlpBlock_0 │ MlpBlock │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_22/MlpBlock_0/Dense_0 │ Dense │ bfloat16[4,64,1024] │ bfloat16[4,64,4096] │ bias: float32[4096] │ +│ │ │ │ │ kernel: float32[1024,4096] │ +│ │ │ │ │ │ +│ │ │ │ │ 4,198,400 (16.8 MB) │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_22/MlpBlock_0/Dropout_0 │ Dropout │ bfloat16[4,64,4096] │ bfloat16[4,64,4096] │ │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_22/MlpBlock_0/Dense_1 │ Dense │ bfloat16[4,64,4096] │ bfloat16[4,64,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[4096,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 4,195,328 (16.8 MB) │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_22/MlpBlock_0/Dropout_1 │ Dropout │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_23 │ DiTBlock │ - bfloat16[4,64,1024] │ bfloat16[4,64,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,64,1024] │ bfloat16[4,64,1024] │ │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_23/Dense_1 │ Dense │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_23/Dense_2 │ Dense │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_23/Dense_3 │ Dense │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼───────────��───────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_23/Dense_4 │ Dense │ float32[4,64,1024] │ bfloat16[4,64,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[1024,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 1,049,600 (4.2 MB) │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_23/LayerNorm_1 │ LayerNorm │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_23/MlpBlock_0 │ MlpBlock │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_23/MlpBlock_0/Dense_0 │ Dense │ bfloat16[4,64,1024] │ bfloat16[4,64,4096] │ bias: float32[4096] │ +│ │ │ │ │ kernel: float32[1024,4096] │ +│ │ │ │ │ │ +│ │ │ │ │ 4,198,400 (16.8 MB) │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_23/MlpBlock_0/Dropout_0 │ Dropout │ bfloat16[4,64,4096] │ bfloat16[4,64,4096] │ │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_23/MlpBlock_0/Dense_1 │ Dense │ bfloat16[4,64,4096] │ bfloat16[4,64,1024] │ bias: float32[1024] │ +│ │ │ │ │ kernel: float32[4096,1024] │ +│ │ │ │ │ │ +│ │ │ │ │ 4,195,328 (16.8 MB) │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ DiTBlock_23/MlpBlock_0/Dropout_1 │ Dropout │ bfloat16[4,64,1024] │ bfloat16[4,64,1024] │ │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ FinalLayer_0 │ FinalLayer │ - bfloat16[4,64,1024] │ bfloat16[4,64,64] │ │ +│ │ │ - 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,64,1024] │ bfloat16[4,64,1024] │ │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ FinalLayer_0/Dense_1 │ Dense │ bfloat16[4,64,1024] │ bfloat16[4,64,64] │ bias: float32[64] │ +│ │ │ │ │ kernel: float32[1024,64] │ +│ │ │ │ │ │ +│ │ │ │ │ 65,600 (262.4 KB) │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│ Embed_0 │ Embed │ int32[1] │ float32[1,1] │ embedding: float32[256,1] │ +│ │ │ │ │ │ +│ │ │ │ │ 256 (1.0 KB) │ +├──────────────────────────────────┼──────────────────┼───────────────────────┼──────────────────────┼───────────────────────────────┤ +│   │   │   │  Total │ 459,235,648 (1.8 GB)  │ +└──────────────────────────────────┴──────────────────┴───────────────────────┴──────────────────────┴───────────────────────────────┘ +  + Total Parameters: 459,235,648 (1.8 GB)  + + +DiT: Input of shape (4, 16, 16, 16) dtype float32 +DiT: After patch embed, shape is (4, 64, 1024) dtype bfloat16 +DiT: Patch Embed of shape (4, 64, 1024) dtype bfloat16 +DiT: Conditioning of shape (1, 1024) dtype float32 +Loaded checkpoint from 15499 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, 64) +('FinalLayer_0', 'Dense_1', 'kernel'): (1, 1024, 64) +('LabelEmbedder_0', 'Embed_0', 'embedding'): (1, 1001, 1024) +('PatchEmbed_0', 'Conv_0', 'bias'): (1, 1024) +('PatchEmbed_0', 'Conv_0', 'kernel'): (1, 2, 2, 16, 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, 128) +doing downsample +Encoder layer (128, 32, 32, 256) +doing downsample +Encoder layer (128, 16, 16, 256) +Encoder layer (128, 16, 16, 512) +Encoder layer final (128, 16, 16, 512) +Encoder layer final (128, 16, 16, 512) +Final embeddings are size (128, 16, 16, 16) +After quant (128, 16, 16, 16) +Calc FID for CFG 1.0 and denoise_timesteps 128 +DiT: Input of shape (512, 16, 16, 16) dtype float32 +DiT: After patch embed, shape is (512, 64, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 64, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +z_vectors shape (128, 16, 16, 16) +Decoder incoming shape (128, 16, 16, 16) +Decoder input (128, 16, 16, 512) +Mid Block Decoder layer (128, 16, 16, 512) +Mid Block Decoder layer (128, 16, 16, 512) +Decoder layer (128, 32, 32, 512) +Decoder layer (128, 64, 64, 256) +Decoder layer (128, 128, 128, 256) +Decoder layer (128, 256, 256, 128) +Decoder layer (128, 256, 256, 128) +FID is 20.55461883544922 +(512, 256, 256, 3) +Calc FID for CFG 1.0 and denoise_timesteps 64 +DiT: Input of shape (512, 16, 16, 16) dtype float32 +DiT: After patch embed, shape is (512, 64, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 64, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 20.785968780517578 +(512, 256, 256, 3) +Calc FID for CFG 1.0 and denoise_timesteps 32 +DiT: Input of shape (512, 16, 16, 16) dtype float32 +DiT: After patch embed, shape is (512, 64, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 64, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 21.26744842529297 +(512, 256, 256, 3) +Calc FID for CFG 1.0 and denoise_timesteps 16 +DiT: Input of shape (512, 16, 16, 16) dtype float32 +DiT: After patch embed, shape is (512, 64, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 64, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 22.38835334777832 +(512, 256, 256, 3) +Calc FID for CFG 1.0 and denoise_timesteps 8 +DiT: Input of shape (512, 16, 16, 16) dtype float32 +DiT: After patch embed, shape is (512, 64, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 64, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 25.291362762451172 +(512, 256, 256, 3) +Calc FID for CFG 1.0 and denoise_timesteps 4 +DiT: Input of shape (512, 16, 16, 16) dtype float32 +DiT: After patch embed, shape is (512, 64, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 64, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 33.44581604003906 +(512, 256, 256, 3) +Calc FID for CFG 1.0 and denoise_timesteps 2 +DiT: Input of shape (512, 16, 16, 16) dtype float32 +DiT: After patch embed, shape is (512, 64, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 64, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 73.94938659667969 +(512, 256, 256, 3) +Calc FID for CFG 1.0 and denoise_timesteps 1 +DiT: Input of shape (512, 16, 16, 16) dtype float32 +DiT: After patch embed, shape is (512, 64, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 64, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 216.96063232421875 +(512, 256, 256, 3) +Calc FID for CFG 1.25 and denoise_timesteps 128 +DiT: Input of shape (512, 16, 16, 16) dtype float32 +DiT: After patch embed, shape is (512, 64, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 64, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 10.432090759277344 +(512, 256, 256, 3) +Calc FID for CFG 1.25 and denoise_timesteps 64 +DiT: Input of shape (512, 16, 16, 16) dtype float32 +DiT: After patch embed, shape is (512, 64, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 64, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 10.581289291381836 +(512, 256, 256, 3) +Calc FID for CFG 1.25 and denoise_timesteps 32 +DiT: Input of shape (512, 16, 16, 16) dtype float32 +DiT: After patch embed, shape is (512, 64, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 64, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 10.897796630859375 +(512, 256, 256, 3) +Calc FID for CFG 1.25 and denoise_timesteps 16 +DiT: Input of shape (512, 16, 16, 16) dtype float32 +DiT: After patch embed, shape is (512, 64, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 64, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 11.704723358154297 +(512, 256, 256, 3) +Calc FID for CFG 1.25 and denoise_timesteps 8 +DiT: Input of shape (512, 16, 16, 16) dtype float32 +DiT: After patch embed, shape is (512, 64, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 64, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 13.877209663391113 +(512, 256, 256, 3) +Calc FID for CFG 1.25 and denoise_timesteps 4 +DiT: Input of shape (512, 16, 16, 16) dtype float32 +DiT: After patch embed, shape is (512, 64, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 64, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 20.557266235351562 +(512, 256, 256, 3) +Calc FID for CFG 1.25 and denoise_timesteps 2 +DiT: Input of shape (512, 16, 16, 16) dtype float32 +DiT: After patch embed, shape is (512, 64, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 64, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 56.92554473876953 +(512, 256, 256, 3) +Calc FID for CFG 1.25 and denoise_timesteps 1 +DiT: Input of shape (512, 16, 16, 16) dtype float32 +DiT: After patch embed, shape is (512, 64, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 64, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 209.0152587890625 +(512, 256, 256, 3) +Calc FID for CFG 1.5 and denoise_timesteps 128 +DiT: Input of shape (512, 16, 16, 16) dtype float32 +DiT: After patch embed, shape is (512, 64, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 64, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 6.526396751403809 +(512, 256, 256, 3) +Calc FID for CFG 1.5 and denoise_timesteps 64 +DiT: Input of shape (512, 16, 16, 16) dtype float32 +DiT: After patch embed, shape is (512, 64, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 64, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 6.636597633361816 +(512, 256, 256, 3) +Calc FID for CFG 1.5 and denoise_timesteps 32 +DiT: Input of shape (512, 16, 16, 16) dtype float32 +DiT: After patch embed, shape is (512, 64, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 64, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 6.8818182945251465 +(512, 256, 256, 3) +Calc FID for CFG 1.5 and denoise_timesteps 16 +DiT: Input of shape (512, 16, 16, 16) dtype float32 +DiT: After patch embed, shape is (512, 64, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 64, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 7.460043430328369 +(512, 256, 256, 3) +Calc FID for CFG 1.5 and denoise_timesteps 8 +DiT: Input of shape (512, 16, 16, 16) dtype float32 +DiT: After patch embed, shape is (512, 64, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 64, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 9.108461380004883 +(512, 256, 256, 3) +Calc FID for CFG 1.5 and denoise_timesteps 4 +DiT: Input of shape (512, 16, 16, 16) dtype float32 +DiT: After patch embed, shape is (512, 64, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 64, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 14.45309066772461 +(512, 256, 256, 3) +Calc FID for CFG 1.5 and denoise_timesteps 2 +DiT: Input of shape (512, 16, 16, 16) dtype float32 +DiT: After patch embed, shape is (512, 64, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 64, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 48.57585144042969 +(512, 256, 256, 3) +Calc FID for CFG 1.5 and denoise_timesteps 1 +DiT: Input of shape (512, 16, 16, 16) dtype float32 +DiT: After patch embed, shape is (512, 64, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 64, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 203.16036987304688 +(512, 256, 256, 3) +Calc FID for CFG 1.75 and denoise_timesteps 128 +DiT: Input of shape (512, 16, 16, 16) dtype float32 +DiT: After patch embed, shape is (512, 64, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 64, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 5.855515956878662 +(512, 256, 256, 3) +Calc FID for CFG 1.75 and denoise_timesteps 64 +DiT: Input of shape (512, 16, 16, 16) dtype float32 +DiT: After patch embed, shape is (512, 64, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 64, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 5.938434600830078 +(512, 256, 256, 3) +Calc FID for CFG 1.75 and denoise_timesteps 32 +DiT: Input of shape (512, 16, 16, 16) dtype float32 +DiT: After patch embed, shape is (512, 64, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 64, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 6.1262946128845215 +(512, 256, 256, 3) +Calc FID for CFG 1.75 and denoise_timesteps 16 +DiT: Input of shape (512, 16, 16, 16) dtype float32 +DiT: After patch embed, shape is (512, 64, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 64, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 6.5858354568481445 +(512, 256, 256, 3) +Calc FID for CFG 1.75 and denoise_timesteps 8 +DiT: Input of shape (512, 16, 16, 16) dtype float32 +DiT: After patch embed, shape is (512, 64, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 64, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 7.889256000518799 +(512, 256, 256, 3) +Calc FID for CFG 1.75 and denoise_timesteps 4 +DiT: Input of shape (512, 16, 16, 16) dtype float32 +DiT: After patch embed, shape is (512, 64, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 64, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 12.243061065673828 +(512, 256, 256, 3) +Calc FID for CFG 1.75 and denoise_timesteps 2 +DiT: Input of shape (512, 16, 16, 16) dtype float32 +DiT: After patch embed, shape is (512, 64, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 64, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 46.02006530761719 +(512, 256, 256, 3) +Calc FID for CFG 1.75 and denoise_timesteps 1 +DiT: Input of shape (512, 16, 16, 16) dtype float32 +DiT: After patch embed, shape is (512, 64, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 64, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 198.68731689453125 +(512, 256, 256, 3) +Calc FID for CFG 2.0 and denoise_timesteps 128 +DiT: Input of shape (512, 16, 16, 16) dtype float32 +DiT: After patch embed, shape is (512, 64, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 64, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 6.705183982849121 +(512, 256, 256, 3) +Calc FID for CFG 2.0 and denoise_timesteps 64 +DiT: Input of shape (512, 16, 16, 16) dtype float32 +DiT: After patch embed, shape is (512, 64, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 64, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 6.767923355102539 +(512, 256, 256, 3) +Calc FID for CFG 2.0 and denoise_timesteps 32 +DiT: Input of shape (512, 16, 16, 16) dtype float32 +DiT: After patch embed, shape is (512, 64, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 64, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 6.9428558349609375 +(512, 256, 256, 3) +Calc FID for CFG 2.0 and denoise_timesteps 16 +DiT: Input of shape (512, 16, 16, 16) dtype float32 +DiT: After patch embed, shape is (512, 64, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 64, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 7.327467441558838 +(512, 256, 256, 3) +Calc FID for CFG 2.0 and denoise_timesteps 8 +DiT: Input of shape (512, 16, 16, 16) dtype float32 +DiT: After patch embed, shape is (512, 64, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 64, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 8.378372192382812 +(512, 256, 256, 3) +Calc FID for CFG 2.0 and denoise_timesteps 4 +DiT: Input of shape (512, 16, 16, 16) dtype float32 +DiT: After patch embed, shape is (512, 64, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 64, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 11.983981132507324 +(512, 256, 256, 3) +Calc FID for CFG 2.0 and denoise_timesteps 2 +DiT: Input of shape (512, 16, 16, 16) dtype float32 +DiT: After patch embed, shape is (512, 64, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 64, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 47.12152862548828 +(512, 256, 256, 3) +Calc FID for CFG 2.0 and denoise_timesteps 1 +DiT: Input of shape (512, 16, 16, 16) dtype float32 +DiT: After patch embed, shape is (512, 64, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 64, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 195.67784118652344 +(512, 256, 256, 3) +Calc FID for CFG 2.25 and denoise_timesteps 128 +DiT: Input of shape (512, 16, 16, 16) dtype float32 +DiT: After patch embed, shape is (512, 64, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 64, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 8.251200675964355 +(512, 256, 256, 3) +Calc FID for CFG 2.25 and denoise_timesteps 64 +DiT: Input of shape (512, 16, 16, 16) dtype float32 +DiT: After patch embed, shape is (512, 64, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 64, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 8.30540943145752 +(512, 256, 256, 3) +Calc FID for CFG 2.25 and denoise_timesteps 32 +DiT: Input of shape (512, 16, 16, 16) dtype float32 +DiT: After patch embed, shape is (512, 64, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 64, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 8.445191383361816 +(512, 256, 256, 3) +Calc FID for CFG 2.25 and denoise_timesteps 16 +DiT: Input of shape (512, 16, 16, 16) dtype float32 +DiT: After patch embed, shape is (512, 64, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 64, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 8.751694679260254 +(512, 256, 256, 3) +Calc FID for CFG 2.25 and denoise_timesteps 8 +DiT: Input of shape (512, 16, 16, 16) dtype float32 +DiT: After patch embed, shape is (512, 64, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 64, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 9.606956481933594 +(512, 256, 256, 3) +Calc FID for CFG 2.25 and denoise_timesteps 4 +DiT: Input of shape (512, 16, 16, 16) dtype float32 +DiT: After patch embed, shape is (512, 64, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 64, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 12.645600318908691 +(512, 256, 256, 3) +Calc FID for CFG 2.25 and denoise_timesteps 2 +DiT: Input of shape (512, 16, 16, 16) dtype float32 +DiT: After patch embed, shape is (512, 64, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 64, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 50.41095733642578 +(512, 256, 256, 3) +Calc FID for CFG 2.25 and denoise_timesteps 1 +DiT: Input of shape (512, 16, 16, 16) dtype float32 +DiT: After patch embed, shape is (512, 64, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 64, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 193.6466064453125 +(512, 256, 256, 3) +Calc FID for CFG 2.5 and denoise_timesteps 128 +DiT: Input of shape (512, 16, 16, 16) dtype float32 +DiT: After patch embed, shape is (512, 64, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 64, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 9.968548774719238 +(512, 256, 256, 3) +Calc FID for CFG 2.5 and denoise_timesteps 64 +DiT: Input of shape (512, 16, 16, 16) dtype float32 +DiT: After patch embed, shape is (512, 64, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 64, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 10.04426383972168 +(512, 256, 256, 3) +Calc FID for CFG 2.5 and denoise_timesteps 32 +DiT: Input of shape (512, 16, 16, 16) dtype float32 +DiT: After patch embed, shape is (512, 64, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 64, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 10.183393478393555 +(512, 256, 256, 3) +Calc FID for CFG 2.5 and denoise_timesteps 16 +DiT: Input of shape (512, 16, 16, 16) dtype float32 +DiT: After patch embed, shape is (512, 64, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 64, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 10.4303560256958 +(512, 256, 256, 3) +Calc FID for CFG 2.5 and denoise_timesteps 8 +DiT: Input of shape (512, 16, 16, 16) dtype float32 +DiT: After patch embed, shape is (512, 64, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 64, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 11.132194519042969 +(512, 256, 256, 3) +Calc FID for CFG 2.5 and denoise_timesteps 4 +DiT: Input of shape (512, 16, 16, 16) dtype float32 +DiT: After patch embed, shape is (512, 64, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 64, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 13.655767440795898 +(512, 256, 256, 3) +Calc FID for CFG 2.5 and denoise_timesteps 2 +DiT: Input of shape (512, 16, 16, 16) dtype float32 +DiT: After patch embed, shape is (512, 64, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 64, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 55.19169616699219 +(512, 256, 256, 3) +Calc FID for CFG 2.5 and denoise_timesteps 1 +DiT: Input of shape (512, 16, 16, 16) dtype float32 +DiT: After patch embed, shape is (512, 64, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 64, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 192.20343017578125 +(512, 256, 256, 3) +Calc FID for CFG 2.75 and denoise_timesteps 128 +DiT: Input of shape (512, 16, 16, 16) dtype float32 +DiT: After patch embed, shape is (512, 64, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 64, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 11.731626510620117 +(512, 256, 256, 3) +Calc FID for CFG 2.75 and denoise_timesteps 64 +DiT: Input of shape (512, 16, 16, 16) dtype float32 +DiT: After patch embed, shape is (512, 64, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 64, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 11.788705825805664 +(512, 256, 256, 3) +Calc FID for CFG 2.75 and denoise_timesteps 32 +DiT: Input of shape (512, 16, 16, 16) dtype float32 +DiT: After patch embed, shape is (512, 64, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 64, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 11.926319122314453 +(512, 256, 256, 3) +Calc FID for CFG 2.75 and denoise_timesteps 16 +DiT: Input of shape (512, 16, 16, 16) dtype float32 +DiT: After patch embed, shape is (512, 64, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 64, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 12.144671440124512 +(512, 256, 256, 3) +Calc FID for CFG 2.75 and denoise_timesteps 8 +DiT: Input of shape (512, 16, 16, 16) dtype float32 +DiT: After patch embed, shape is (512, 64, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 64, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 12.685546875 +(512, 256, 256, 3) +Calc FID for CFG 2.75 and denoise_timesteps 4 +DiT: Input of shape (512, 16, 16, 16) dtype float32 +DiT: After patch embed, shape is (512, 64, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 64, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 14.735636711120605 +(512, 256, 256, 3) +Calc FID for CFG 2.75 and denoise_timesteps 2 +DiT: Input of shape (512, 16, 16, 16) dtype float32 +DiT: After patch embed, shape is (512, 64, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 64, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 60.91100311279297 +(512, 256, 256, 3) +Calc FID for CFG 2.75 and denoise_timesteps 1 +DiT: Input of shape (512, 16, 16, 16) dtype float32 +DiT: After patch embed, shape is (512, 64, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 64, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 191.12744140625 +(512, 256, 256, 3) +Calc FID for CFG 3.0 and denoise_timesteps 128 +DiT: Input of shape (512, 16, 16, 16) dtype float32 +DiT: After patch embed, shape is (512, 64, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 64, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 13.366455078125 +(512, 256, 256, 3) +Calc FID for CFG 3.0 and denoise_timesteps 64 +DiT: Input of shape (512, 16, 16, 16) dtype float32 +DiT: After patch embed, shape is (512, 64, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 64, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 13.395753860473633 +(512, 256, 256, 3) +Calc FID for CFG 3.0 and denoise_timesteps 32 +DiT: Input of shape (512, 16, 16, 16) dtype float32 +DiT: After patch embed, shape is (512, 64, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 64, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 13.500991821289062 +(512, 256, 256, 3) +Calc FID for CFG 3.0 and denoise_timesteps 16 +DiT: Input of shape (512, 16, 16, 16) dtype float32 +DiT: After patch embed, shape is (512, 64, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 64, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 13.753835678100586 +(512, 256, 256, 3) +Calc FID for CFG 3.0 and denoise_timesteps 8 +DiT: Input of shape (512, 16, 16, 16) dtype float32 +DiT: After patch embed, shape is (512, 64, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 64, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 14.131613731384277 +(512, 256, 256, 3) +Calc FID for CFG 3.0 and denoise_timesteps 4 +DiT: Input of shape (512, 16, 16, 16) dtype float32 +DiT: After patch embed, shape is (512, 64, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 64, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 15.804227828979492 +(512, 256, 256, 3) +Calc FID for CFG 3.0 and denoise_timesteps 2 +DiT: Input of shape (512, 16, 16, 16) dtype float32 +DiT: After patch embed, shape is (512, 64, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 64, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 67.34161376953125 +(512, 256, 256, 3) +Calc FID for CFG 3.0 and denoise_timesteps 1 +DiT: Input of shape (512, 16, 16, 16) dtype float32 +DiT: After patch embed, shape is (512, 64, 1024) dtype bfloat16 +DiT: Patch Embed of shape (512, 64, 1024) dtype bfloat16 +DiT: Conditioning of shape (512, 1024) dtype float32 +FID is 190.349365234375