Upload folder using huggingface_hub

2e-5_no_sampling/log.txt

@@ -67,7 +67,7 @@ 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 16591587 seconds ago.
+Loaded checkpoint from 16697582 seconds ago.
 Loaded model with step 511001
 ┌──────────────────────────────────────────────────────────────────────────────┐
 │                                    TPU 0                                     │
@@ -768,3 +768,1671 @@ DiT: Input of shape (4, 32, 32, 4) dtype float32
 DiT: After patch embed, shape is (4, 256, 768) dtype bfloat16
 DiT: Patch Embed of shape (4, 256, 768) dtype bfloat16
 DiT: Conditioning of shape (1, 768) dtype float32
+Loaded checkpoint from 153447 seconds ago.
+
+ parameter shapes:
+('PatchEmbed_0', 'Conv_0', 'kernel'): (2, 2, 4, 768)
+('PatchEmbed_0', 'Conv_0', 'bias'): (768,)
+('TimestepEmbedder_0', 'Dense_0', 'kernel'): (256, 768)
+('TimestepEmbedder_0', 'Dense_0', 'bias'): (768,)
+('TimestepEmbedder_0', 'Dense_1', 'kernel'): (768, 768)
+('TimestepEmbedder_0', 'Dense_1', 'bias'): (768,)
+('TimestepEmbedder_1', 'Dense_0', 'kernel'): (256, 768)
+('TimestepEmbedder_1', 'Dense_0', 'bias'): (768,)
+('TimestepEmbedder_1', 'Dense_1', 'kernel'): (768, 768)
+('TimestepEmbedder_1', 'Dense_1', 'bias'): (768,)
+('LabelEmbedder_0', 'Embed_0', 'embedding'): (1001, 768)
+('DiTBlock_0', 'Dense_0', 'kernel'): (768, 4608)
+('DiTBlock_0', 'Dense_0', 'bias'): (4608,)
+('DiTBlock_0', 'Dense_1', 'kernel'): (768, 768)
+('DiTBlock_0', 'Dense_1', 'bias'): (768,)
+('DiTBlock_0', 'Dense_2', 'kernel'): (768, 768)
+('DiTBlock_0', 'Dense_2', 'bias'): (768,)
+('DiTBlock_0', 'Dense_3', 'kernel'): (768, 768)
+('DiTBlock_0', 'Dense_3', 'bias'): (768,)
+('DiTBlock_0', 'Dense_4', 'kernel'): (768, 768)
+('DiTBlock_0', 'Dense_4', 'bias'): (768,)
+('DiTBlock_0', 'MlpBlock_0', 'Dense_0', 'kernel'): (768, 3072)
+('DiTBlock_0', 'MlpBlock_0', 'Dense_0', 'bias'): (3072,)
+('DiTBlock_0', 'MlpBlock_0', 'Dense_1', 'kernel'): (3072, 768)
+('DiTBlock_0', 'MlpBlock_0', 'Dense_1', 'bias'): (768,)
+('DiTBlock_1', 'Dense_0', 'kernel'): (768, 4608)
+('DiTBlock_1', 'Dense_0', 'bias'): (4608,)
+('DiTBlock_1', 'Dense_1', 'kernel'): (768, 768)
+('DiTBlock_1', 'Dense_1', 'bias'): (768,)
+('DiTBlock_1', 'Dense_2', 'kernel'): (768, 768)
+('DiTBlock_1', 'Dense_2', 'bias'): (768,)
+('DiTBlock_1', 'Dense_3', 'kernel'): (768, 768)
+('DiTBlock_1', 'Dense_3', 'bias'): (768,)
+('DiTBlock_1', 'Dense_4', 'kernel'): (768, 768)
+('DiTBlock_1', 'Dense_4', 'bias'): (768,)
+('DiTBlock_1', 'MlpBlock_0', 'Dense_0', 'kernel'): (768, 3072)
+('DiTBlock_1', 'MlpBlock_0', 'Dense_0', 'bias'): (3072,)
+('DiTBlock_1', 'MlpBlock_0', 'Dense_1', 'kernel'): (3072, 768)
+('DiTBlock_1', 'MlpBlock_0', 'Dense_1', 'bias'): (768,)
+('DiTBlock_2', 'Dense_0', 'kernel'): (768, 4608)
+('DiTBlock_2', 'Dense_0', 'bias'): (4608,)
+('DiTBlock_2', 'Dense_1', 'kernel'): (768, 768)
+('DiTBlock_2', 'Dense_1', 'bias'): (768,)
+('DiTBlock_2', 'Dense_2', 'kernel'): (768, 768)
+('DiTBlock_2', 'Dense_2', 'bias'): (768,)
+('DiTBlock_2', 'Dense_3', 'kernel'): (768, 768)
+('DiTBlock_2', 'Dense_3', 'bias'): (768,)
+('DiTBlock_2', 'Dense_4', 'kernel'): (768, 768)
+('DiTBlock_2', 'Dense_4', 'bias'): (768,)
+('DiTBlock_2', 'MlpBlock_0', 'Dense_0', 'kernel'): (768, 3072)
+('DiTBlock_2', 'MlpBlock_0', 'Dense_0', 'bias'): (3072,)
+('DiTBlock_2', 'MlpBlock_0', 'Dense_1', 'kernel'): (3072, 768)
+('DiTBlock_2', 'MlpBlock_0', 'Dense_1', 'bias'): (768,)
+('DiTBlock_3', 'Dense_0', 'kernel'): (768, 4608)
+('DiTBlock_3', 'Dense_0', 'bias'): (4608,)
+('DiTBlock_3', 'Dense_1', 'kernel'): (768, 768)
+('DiTBlock_3', 'Dense_1', 'bias'): (768,)
+('DiTBlock_3', 'Dense_2', 'kernel'): (768, 768)
+('DiTBlock_3', 'Dense_2', 'bias'): (768,)
+('DiTBlock_3', 'Dense_3', 'kernel'): (768, 768)
+('DiTBlock_3', 'Dense_3', 'bias'): (768,)
+('DiTBlock_3', 'Dense_4', 'kernel'): (768, 768)
+('DiTBlock_3', 'Dense_4', 'bias'): (768,)
+('DiTBlock_3', 'MlpBlock_0', 'Dense_0', 'kernel'): (768, 3072)
+('DiTBlock_3', 'MlpBlock_0', 'Dense_0', 'bias'): (3072,)
+('DiTBlock_3', 'MlpBlock_0', 'Dense_1', 'kernel'): (3072, 768)
+('DiTBlock_3', 'MlpBlock_0', 'Dense_1', 'bias'): (768,)
+('DiTBlock_4', 'Dense_0', 'kernel'): (768, 4608)
+('DiTBlock_4', 'Dense_0', 'bias'): (4608,)
+('DiTBlock_4', 'Dense_1', 'kernel'): (768, 768)
+('DiTBlock_4', 'Dense_1', 'bias'): (768,)
+('DiTBlock_4', 'Dense_2', 'kernel'): (768, 768)
+('DiTBlock_4', 'Dense_2', 'bias'): (768,)
+('DiTBlock_4', 'Dense_3', 'kernel'): (768, 768)
+('DiTBlock_4', 'Dense_3', 'bias'): (768,)
+('DiTBlock_4', 'Dense_4', 'kernel'): (768, 768)
+('DiTBlock_4', 'Dense_4', 'bias'): (768,)
+('DiTBlock_4', 'MlpBlock_0', 'Dense_0', 'kernel'): (768, 3072)
+('DiTBlock_4', 'MlpBlock_0', 'Dense_0', 'bias'): (3072,)
+('DiTBlock_4', 'MlpBlock_0', 'Dense_1', 'kernel'): (3072, 768)
+('DiTBlock_4', 'MlpBlock_0', 'Dense_1', 'bias'): (768,)
+('DiTBlock_5', 'Dense_0', 'kernel'): (768, 4608)
+('DiTBlock_5', 'Dense_0', 'bias'): (4608,)
+('DiTBlock_5', 'Dense_1', 'kernel'): (768, 768)
+('DiTBlock_5', 'Dense_1', 'bias'): (768,)
+('DiTBlock_5', 'Dense_2', 'kernel'): (768, 768)
+('DiTBlock_5', 'Dense_2', 'bias'): (768,)
+('DiTBlock_5', 'Dense_3', 'kernel'): (768, 768)
+('DiTBlock_5', 'Dense_3', 'bias'): (768,)
+('DiTBlock_5', 'Dense_4', 'kernel'): (768, 768)
+('DiTBlock_5', 'Dense_4', 'bias'): (768,)
+('DiTBlock_5', 'MlpBlock_0', 'Dense_0', 'kernel'): (768, 3072)
+('DiTBlock_5', 'MlpBlock_0', 'Dense_0', 'bias'): (3072,)
+('DiTBlock_5', 'MlpBlock_0', 'Dense_1', 'kernel'): (3072, 768)
+('DiTBlock_5', 'MlpBlock_0', 'Dense_1', 'bias'): (768,)
+('DiTBlock_6', 'Dense_0', 'kernel'): (768, 4608)
+('DiTBlock_6', 'Dense_0', 'bias'): (4608,)
+('DiTBlock_6', 'Dense_1', 'kernel'): (768, 768)
+('DiTBlock_6', 'Dense_1', 'bias'): (768,)
+('DiTBlock_6', 'Dense_2', 'kernel'): (768, 768)
+('DiTBlock_6', 'Dense_2', 'bias'): (768,)
+('DiTBlock_6', 'Dense_3', 'kernel'): (768, 768)
+('DiTBlock_6', 'Dense_3', 'bias'): (768,)
+('DiTBlock_6', 'Dense_4', 'kernel'): (768, 768)
+('DiTBlock_6', 'Dense_4', 'bias'): (768,)
+('DiTBlock_6', 'MlpBlock_0', 'Dense_0', 'kernel'): (768, 3072)
+('DiTBlock_6', 'MlpBlock_0', 'Dense_0', 'bias'): (3072,)
+('DiTBlock_6', 'MlpBlock_0', 'Dense_1', 'kernel'): (3072, 768)
+('DiTBlock_6', 'MlpBlock_0', 'Dense_1', 'bias'): (768,)
+('DiTBlock_7', 'Dense_0', 'kernel'): (768, 4608)
+('DiTBlock_7', 'Dense_0', 'bias'): (4608,)
+('DiTBlock_7', 'Dense_1', 'kernel'): (768, 768)
+('DiTBlock_7', 'Dense_1', 'bias'): (768,)
+('DiTBlock_7', 'Dense_2', 'kernel'): (768, 768)
+('DiTBlock_7', 'Dense_2', 'bias'): (768,)
+('DiTBlock_7', 'Dense_3', 'kernel'): (768, 768)
+('DiTBlock_7', 'Dense_3', 'bias'): (768,)
+('DiTBlock_7', 'Dense_4', 'kernel'): (768, 768)
+('DiTBlock_7', 'Dense_4', 'bias'): (768,)
+('DiTBlock_7', 'MlpBlock_0', 'Dense_0', 'kernel'): (768, 3072)
+('DiTBlock_7', 'MlpBlock_0', 'Dense_0', 'bias'): (3072,)
+('DiTBlock_7', 'MlpBlock_0', 'Dense_1', 'kernel'): (3072, 768)
+('DiTBlock_7', 'MlpBlock_0', 'Dense_1', 'bias'): (768,)
+('DiTBlock_8', 'Dense_0', 'kernel'): (768, 4608)
+('DiTBlock_8', 'Dense_0', 'bias'): (4608,)
+('DiTBlock_8', 'Dense_1', 'kernel'): (768, 768)
+('DiTBlock_8', 'Dense_1', 'bias'): (768,)
+('DiTBlock_8', 'Dense_2', 'kernel'): (768, 768)
+('DiTBlock_8', 'Dense_2', 'bias'): (768,)
+('DiTBlock_8', 'Dense_3', 'kernel'): (768, 768)
+('DiTBlock_8', 'Dense_3', 'bias'): (768,)
+('DiTBlock_8', 'Dense_4', 'kernel'): (768, 768)
+('DiTBlock_8', 'Dense_4', 'bias'): (768,)
+('DiTBlock_8', 'MlpBlock_0', 'Dense_0', 'kernel'): (768, 3072)
+('DiTBlock_8', 'MlpBlock_0', 'Dense_0', 'bias'): (3072,)
+('DiTBlock_8', 'MlpBlock_0', 'Dense_1', 'kernel'): (3072, 768)
+('DiTBlock_8', 'MlpBlock_0', 'Dense_1', 'bias'): (768,)
+('DiTBlock_9', 'Dense_0', 'kernel'): (768, 4608)
+('DiTBlock_9', 'Dense_0', 'bias'): (4608,)
+('DiTBlock_9', 'Dense_1', 'kernel'): (768, 768)
+('DiTBlock_9', 'Dense_1', 'bias'): (768,)
+('DiTBlock_9', 'Dense_2', 'kernel'): (768, 768)
+('DiTBlock_9', 'Dense_2', 'bias'): (768,)
+('DiTBlock_9', 'Dense_3', 'kernel'): (768, 768)
+('DiTBlock_9', 'Dense_3', 'bias'): (768,)
+('DiTBlock_9', 'Dense_4', 'kernel'): (768, 768)
+('DiTBlock_9', 'Dense_4', 'bias'): (768,)
+('DiTBlock_9', 'MlpBlock_0', 'Dense_0', 'kernel'): (768, 3072)
+('DiTBlock_9', 'MlpBlock_0', 'Dense_0', 'bias'): (3072,)
+('DiTBlock_9', 'MlpBlock_0', 'Dense_1', 'kernel'): (3072, 768)
+('DiTBlock_9', 'MlpBlock_0', 'Dense_1', 'bias'): (768,)
+('DiTBlock_10', 'Dense_0', 'kernel'): (768, 4608)
+('DiTBlock_10', 'Dense_0', 'bias'): (4608,)
+('DiTBlock_10', 'Dense_1', 'kernel'): (768, 768)
+('DiTBlock_10', 'Dense_1', 'bias'): (768,)
+('DiTBlock_10', 'Dense_2', 'kernel'): (768, 768)
+('DiTBlock_10', 'Dense_2', 'bias'): (768,)
+('DiTBlock_10', 'Dense_3', 'kernel'): (768, 768)
+('DiTBlock_10', 'Dense_3', 'bias'): (768,)
+('DiTBlock_10', 'Dense_4', 'kernel'): (768, 768)
+('DiTBlock_10', 'Dense_4', 'bias'): (768,)
+('DiTBlock_10', 'MlpBlock_0', 'Dense_0', 'kernel'): (768, 3072)
+('DiTBlock_10', 'MlpBlock_0', 'Dense_0', 'bias'): (3072,)
+('DiTBlock_10', 'MlpBlock_0', 'Dense_1', 'kernel'): (3072, 768)
+('DiTBlock_10', 'MlpBlock_0', 'Dense_1', 'bias'): (768,)
+('DiTBlock_11', 'Dense_0', 'kernel'): (768, 4608)
+('DiTBlock_11', 'Dense_0', 'bias'): (4608,)
+('DiTBlock_11', 'Dense_1', 'kernel'): (768, 768)
+('DiTBlock_11', 'Dense_1', 'bias'): (768,)
+('DiTBlock_11', 'Dense_2', 'kernel'): (768, 768)
+('DiTBlock_11', 'Dense_2', 'bias'): (768,)
+('DiTBlock_11', 'Dense_3', 'kernel'): (768, 768)
+('DiTBlock_11', 'Dense_3', 'bias'): (768,)
+('DiTBlock_11', 'Dense_4', 'kernel'): (768, 768)
+('DiTBlock_11', 'Dense_4', 'bias'): (768,)
+('DiTBlock_11', 'MlpBlock_0', 'Dense_0', 'kernel'): (768, 3072)
+('DiTBlock_11', 'MlpBlock_0', 'Dense_0', 'bias'): (3072,)
+('DiTBlock_11', 'MlpBlock_0', 'Dense_1', 'kernel'): (3072, 768)
+('DiTBlock_11', 'MlpBlock_0', 'Dense_1', 'bias'): (768,)
+('FinalLayer_0', 'Dense_0', 'kernel'): (768, 1536)
+('FinalLayer_0', 'Dense_0', 'bias'): (1536,)
+('FinalLayer_0', 'Dense_1', 'kernel'): (768, 16)
+('FinalLayer_0', 'Dense_1', 'bias'): (16,)
+('Embed_0', 'embedding'): (256, 1)
+
+ parameter shapes:
+('DiTBlock_0', 'Dense_0', 'bias'): (1, 4608)
+('DiTBlock_0', 'Dense_0', 'kernel'): (1, 768, 4608)
+('DiTBlock_0', 'Dense_1', 'bias'): (1, 768)
+('DiTBlock_0', 'Dense_1', 'kernel'): (1, 768, 768)
+('DiTBlock_0', 'Dense_2', 'bias'): (1, 768)
+('DiTBlock_0', 'Dense_2', 'kernel'): (1, 768, 768)
+('DiTBlock_0', 'Dense_3', 'bias'): (1, 768)
+('DiTBlock_0', 'Dense_3', 'kernel'): (1, 768, 768)
+('DiTBlock_0', 'Dense_4', 'bias'): (1, 768)
+('DiTBlock_0', 'Dense_4', 'kernel'): (1, 768, 768)
+('DiTBlock_0', 'MlpBlock_0', 'Dense_0', 'bias'): (1, 3072)
+('DiTBlock_0', 'MlpBlock_0', 'Dense_0', 'kernel'): (1, 768, 3072)
+('DiTBlock_0', 'MlpBlock_0', 'Dense_1', 'bias'): (1, 768)
+('DiTBlock_0', 'MlpBlock_0', 'Dense_1', 'kernel'): (1, 3072, 768)
+('DiTBlock_1', 'Dense_0', 'bias'): (1, 4608)
+('DiTBlock_1', 'Dense_0', 'kernel'): (1, 768, 4608)
+('DiTBlock_1', 'Dense_1', 'bias'): (1, 768)
+('DiTBlock_1', 'Dense_1', 'kernel'): (1, 768, 768)
+('DiTBlock_1', 'Dense_2', 'bias'): (1, 768)
+('DiTBlock_1', 'Dense_2', 'kernel'): (1, 768, 768)
+('DiTBlock_1', 'Dense_3', 'bias'): (1, 768)
+('DiTBlock_1', 'Dense_3', 'kernel'): (1, 768, 768)
+('DiTBlock_1', 'Dense_4', 'bias'): (1, 768)
+('DiTBlock_1', 'Dense_4', 'kernel'): (1, 768, 768)
+('DiTBlock_1', 'MlpBlock_0', 'Dense_0', 'bias'): (1, 3072)
+('DiTBlock_1', 'MlpBlock_0', 'Dense_0', 'kernel'): (1, 768, 3072)
+('DiTBlock_1', 'MlpBlock_0', 'Dense_1', 'bias'): (1, 768)
+('DiTBlock_1', 'MlpBlock_0', 'Dense_1', 'kernel'): (1, 3072, 768)
+('DiTBlock_10', 'Dense_0', 'bias'): (1, 4608)
+('DiTBlock_10', 'Dense_0', 'kernel'): (1, 768, 4608)
+('DiTBlock_10', 'Dense_1', 'bias'): (1, 768)
+('DiTBlock_10', 'Dense_1', 'kernel'): (1, 768, 768)
+('DiTBlock_10', 'Dense_2', 'bias'): (1, 768)
+('DiTBlock_10', 'Dense_2', 'kernel'): (1, 768, 768)
+('DiTBlock_10', 'Dense_3', 'bias'): (1, 768)
+('DiTBlock_10', 'Dense_3', 'kernel'): (1, 768, 768)
+('DiTBlock_10', 'Dense_4', 'bias'): (1, 768)
+('DiTBlock_10', 'Dense_4', 'kernel'): (1, 768, 768)
+('DiTBlock_10', 'MlpBlock_0', 'Dense_0', 'bias'): (1, 3072)
+('DiTBlock_10', 'MlpBlock_0', 'Dense_0', 'kernel'): (1, 768, 3072)
+('DiTBlock_10', 'MlpBlock_0', 'Dense_1', 'bias'): (1, 768)
+('DiTBlock_10', 'MlpBlock_0', 'Dense_1', 'kernel'): (1, 3072, 768)
+('DiTBlock_11', 'Dense_0', 'bias'): (1, 4608)
+('DiTBlock_11', 'Dense_0', 'kernel'): (1, 768, 4608)
+('DiTBlock_11', 'Dense_1', 'bias'): (1, 768)
+('DiTBlock_11', 'Dense_1', 'kernel'): (1, 768, 768)
+('DiTBlock_11', 'Dense_2', 'bias'): (1, 768)
+('DiTBlock_11', 'Dense_2', 'kernel'): (1, 768, 768)
+('DiTBlock_11', 'Dense_3', 'bias'): (1, 768)
+('DiTBlock_11', 'Dense_3', 'kernel'): (1, 768, 768)
+('DiTBlock_11', 'Dense_4', 'bias'): (1, 768)
+('DiTBlock_11', 'Dense_4', 'kernel'): (1, 768, 768)
+('DiTBlock_11', 'MlpBlock_0', 'Dense_0', 'bias'): (1, 3072)
+('DiTBlock_11', 'MlpBlock_0', 'Dense_0', 'kernel'): (1, 768, 3072)
+('DiTBlock_11', 'MlpBlock_0', 'Dense_1', 'bias'): (1, 768)
+('DiTBlock_11', 'MlpBlock_0', 'Dense_1', 'kernel'): (1, 3072, 768)
+('DiTBlock_2', 'Dense_0', 'bias'): (1, 4608)
+('DiTBlock_2', 'Dense_0', 'kernel'): (1, 768, 4608)
+('DiTBlock_2', 'Dense_1', 'bias'): (1, 768)
+('DiTBlock_2', 'Dense_1', 'kernel'): (1, 768, 768)
+('DiTBlock_2', 'Dense_2', 'bias'): (1, 768)
+('DiTBlock_2', 'Dense_2', 'kernel'): (1, 768, 768)
+('DiTBlock_2', 'Dense_3', 'bias'): (1, 768)
+('DiTBlock_2', 'Dense_3', 'kernel'): (1, 768, 768)
+('DiTBlock_2', 'Dense_4', 'bias'): (1, 768)
+('DiTBlock_2', 'Dense_4', 'kernel'): (1, 768, 768)
+('DiTBlock_2', 'MlpBlock_0', 'Dense_0', 'bias'): (1, 3072)
+('DiTBlock_2', 'MlpBlock_0', 'Dense_0', 'kernel'): (1, 768, 3072)
+('DiTBlock_2', 'MlpBlock_0', 'Dense_1', 'bias'): (1, 768)
+('DiTBlock_2', 'MlpBlock_0', 'Dense_1', 'kernel'): (1, 3072, 768)
+('DiTBlock_3', 'Dense_0', 'bias'): (1, 4608)
+('DiTBlock_3', 'Dense_0', 'kernel'): (1, 768, 4608)
+('DiTBlock_3', 'Dense_1', 'bias'): (1, 768)
+('DiTBlock_3', 'Dense_1', 'kernel'): (1, 768, 768)
+('DiTBlock_3', 'Dense_2', 'bias'): (1, 768)
+('DiTBlock_3', 'Dense_2', 'kernel'): (1, 768, 768)
+('DiTBlock_3', 'Dense_3', 'bias'): (1, 768)
+('DiTBlock_3', 'Dense_3', 'kernel'): (1, 768, 768)
+('DiTBlock_3', 'Dense_4', 'bias'): (1, 768)
+('DiTBlock_3', 'Dense_4', 'kernel'): (1, 768, 768)
+('DiTBlock_3', 'MlpBlock_0', 'Dense_0', 'bias'): (1, 3072)
+('DiTBlock_3', 'MlpBlock_0', 'Dense_0', 'kernel'): (1, 768, 3072)
+('DiTBlock_3', 'MlpBlock_0', 'Dense_1', 'bias'): (1, 768)
+('DiTBlock_3', 'MlpBlock_0', 'Dense_1', 'kernel'): (1, 3072, 768)
+('DiTBlock_4', 'Dense_0', 'bias'): (1, 4608)
+('DiTBlock_4', 'Dense_0', 'kernel'): (1, 768, 4608)
+('DiTBlock_4', 'Dense_1', 'bias'): (1, 768)
+('DiTBlock_4', 'Dense_1', 'kernel'): (1, 768, 768)
+('DiTBlock_4', 'Dense_2', 'bias'): (1, 768)
+('DiTBlock_4', 'Dense_2', 'kernel'): (1, 768, 768)
+('DiTBlock_4', 'Dense_3', 'bias'): (1, 768)
+('DiTBlock_4', 'Dense_3', 'kernel'): (1, 768, 768)
+('DiTBlock_4', 'Dense_4', 'bias'): (1, 768)
+('DiTBlock_4', 'Dense_4', 'kernel'): (1, 768, 768)
+('DiTBlock_4', 'MlpBlock_0', 'Dense_0', 'bias'): (1, 3072)
+('DiTBlock_4', 'MlpBlock_0', 'Dense_0', 'kernel'): (1, 768, 3072)
+('DiTBlock_4', 'MlpBlock_0', 'Dense_1', 'bias'): (1, 768)
+('DiTBlock_4', 'MlpBlock_0', 'Dense_1', 'kernel'): (1, 3072, 768)
+('DiTBlock_5', 'Dense_0', 'bias'): (1, 4608)
+('DiTBlock_5', 'Dense_0', 'kernel'): (1, 768, 4608)
+('DiTBlock_5', 'Dense_1', 'bias'): (1, 768)
+('DiTBlock_5', 'Dense_1', 'kernel'): (1, 768, 768)
+('DiTBlock_5', 'Dense_2', 'bias'): (1, 768)
+('DiTBlock_5', 'Dense_2', 'kernel'): (1, 768, 768)
+('DiTBlock_5', 'Dense_3', 'bias'): (1, 768)
+('DiTBlock_5', 'Dense_3', 'kernel'): (1, 768, 768)
+('DiTBlock_5', 'Dense_4', 'bias'): (1, 768)
+('DiTBlock_5', 'Dense_4', 'kernel'): (1, 768, 768)
+('DiTBlock_5', 'MlpBlock_0', 'Dense_0', 'bias'): (1, 3072)
+('DiTBlock_5', 'MlpBlock_0', 'Dense_0', 'kernel'): (1, 768, 3072)
+('DiTBlock_5', 'MlpBlock_0', 'Dense_1', 'bias'): (1, 768)
+('DiTBlock_5', 'MlpBlock_0', 'Dense_1', 'kernel'): (1, 3072, 768)
+('DiTBlock_6', 'Dense_0', 'bias'): (1, 4608)
+('DiTBlock_6', 'Dense_0', 'kernel'): (1, 768, 4608)
+('DiTBlock_6', 'Dense_1', 'bias'): (1, 768)
+('DiTBlock_6', 'Dense_1', 'kernel'): (1, 768, 768)
+('DiTBlock_6', 'Dense_2', 'bias'): (1, 768)
+('DiTBlock_6', 'Dense_2', 'kernel'): (1, 768, 768)
+('DiTBlock_6', 'Dense_3', 'bias'): (1, 768)
+('DiTBlock_6', 'Dense_3', 'kernel'): (1, 768, 768)
+('DiTBlock_6', 'Dense_4', 'bias'): (1, 768)
+('DiTBlock_6', 'Dense_4', 'kernel'): (1, 768, 768)
+('DiTBlock_6', 'MlpBlock_0', 'Dense_0', 'bias'): (1, 3072)
+('DiTBlock_6', 'MlpBlock_0', 'Dense_0', 'kernel'): (1, 768, 3072)
+('DiTBlock_6', 'MlpBlock_0', 'Dense_1', 'bias'): (1, 768)
+('DiTBlock_6', 'MlpBlock_0', 'Dense_1', 'kernel'): (1, 3072, 768)
+('DiTBlock_7', 'Dense_0', 'bias'): (1, 4608)
+('DiTBlock_7', 'Dense_0', 'kernel'): (1, 768, 4608)
+('DiTBlock_7', 'Dense_1', 'bias'): (1, 768)
+('DiTBlock_7', 'Dense_1', 'kernel'): (1, 768, 768)
+('DiTBlock_7', 'Dense_2', 'bias'): (1, 768)
+('DiTBlock_7', 'Dense_2', 'kernel'): (1, 768, 768)
+('DiTBlock_7', 'Dense_3', 'bias'): (1, 768)
+('DiTBlock_7', 'Dense_3', 'kernel'): (1, 768, 768)
+('DiTBlock_7', 'Dense_4', 'bias'): (1, 768)
+('DiTBlock_7', 'Dense_4', 'kernel'): (1, 768, 768)
+('DiTBlock_7', 'MlpBlock_0', 'Dense_0', 'bias'): (1, 3072)
+('DiTBlock_7', 'MlpBlock_0', 'Dense_0', 'kernel'): (1, 768, 3072)
+('DiTBlock_7', 'MlpBlock_0', 'Dense_1', 'bias'): (1, 768)
+('DiTBlock_7', 'MlpBlock_0', 'Dense_1', 'kernel'): (1, 3072, 768)
+('DiTBlock_8', 'Dense_0', 'bias'): (1, 4608)
+('DiTBlock_8', 'Dense_0', 'kernel'): (1, 768, 4608)
+('DiTBlock_8', 'Dense_1', 'bias'): (1, 768)
+('DiTBlock_8', 'Dense_1', 'kernel'): (1, 768, 768)
+('DiTBlock_8', 'Dense_2', 'bias'): (1, 768)
+('DiTBlock_8', 'Dense_2', 'kernel'): (1, 768, 768)
+('DiTBlock_8', 'Dense_3', 'bias'): (1, 768)
+('DiTBlock_8', 'Dense_3', 'kernel'): (1, 768, 768)
+('DiTBlock_8', 'Dense_4', 'bias'): (1, 768)
+('DiTBlock_8', 'Dense_4', 'kernel'): (1, 768, 768)
+('DiTBlock_8', 'MlpBlock_0', 'Dense_0', 'bias'): (1, 3072)
+('DiTBlock_8', 'MlpBlock_0', 'Dense_0', 'kernel'): (1, 768, 3072)
+('DiTBlock_8', 'MlpBlock_0', 'Dense_1', 'bias'): (1, 768)
+('DiTBlock_8', 'MlpBlock_0', 'Dense_1', 'kernel'): (1, 3072, 768)
+('DiTBlock_9', 'Dense_0', 'bias'): (1, 4608)
+('DiTBlock_9', 'Dense_0', 'kernel'): (1, 768, 4608)
+('DiTBlock_9', 'Dense_1', 'bias'): (1, 768)
+('DiTBlock_9', 'Dense_1', 'kernel'): (1, 768, 768)
+('DiTBlock_9', 'Dense_2', 'bias'): (1, 768)
+('DiTBlock_9', 'Dense_2', 'kernel'): (1, 768, 768)
+('DiTBlock_9', 'Dense_3', 'bias'): (1, 768)
+('DiTBlock_9', 'Dense_3', 'kernel'): (1, 768, 768)
+('DiTBlock_9', 'Dense_4', 'bias'): (1, 768)
+('DiTBlock_9', 'Dense_4', 'kernel'): (1, 768, 768)
+('DiTBlock_9', 'MlpBlock_0', 'Dense_0', 'bias'): (1, 3072)
+('DiTBlock_9', 'MlpBlock_0', 'Dense_0', 'kernel'): (1, 768, 3072)
+('DiTBlock_9', 'MlpBlock_0', 'Dense_1', 'bias'): (1, 768)
+('DiTBlock_9', 'MlpBlock_0', 'Dense_1', 'kernel'): (1, 3072, 768)
+('Embed_0', 'embedding'): (1, 256, 1)
+('FinalLayer_0', 'Dense_0', 'bias'): (1, 1536)
+('FinalLayer_0', 'Dense_0', 'kernel'): (1, 768, 1536)
+('FinalLayer_0', 'Dense_1', 'bias'): (1, 16)
+('FinalLayer_0', 'Dense_1', 'kernel'): (1, 768, 16)
+('LabelEmbedder_0', 'Embed_0', 'embedding'): (1, 1001, 768)
+('PatchEmbed_0', 'Conv_0', 'bias'): (1, 768)
+('PatchEmbed_0', 'Conv_0', 'kernel'): (1, 2, 2, 4, 768)
+('TimestepEmbedder_0', 'Dense_0', 'bias'): (1, 768)
+('TimestepEmbedder_0', 'Dense_0', 'kernel'): (1, 256, 768)
+('TimestepEmbedder_0', 'Dense_1', 'bias'): (1, 768)
+('TimestepEmbedder_0', 'Dense_1', 'kernel'): (1, 768, 768)
+('TimestepEmbedder_1', 'Dense_0', 'bias'): (1, 768)
+('TimestepEmbedder_1', 'Dense_0', 'kernel'): (1, 256, 768)
+('TimestepEmbedder_1', 'Dense_1', 'bias'): (1, 768)
+('TimestepEmbedder_1', 'Dense_1', 'kernel'): (1, 768, 768)
+
+ parameter shapes:
+('DiTBlock_0', 'Dense_0', 'bias'): (1, 4608)
+('DiTBlock_0', 'Dense_0', 'kernel'): (1, 768, 4608)
+('DiTBlock_0', 'Dense_1', 'bias'): (1, 768)
+('DiTBlock_0', 'Dense_1', 'kernel'): (1, 768, 768)
+('DiTBlock_0', 'Dense_2', 'bias'): (1, 768)
+('DiTBlock_0', 'Dense_2', 'kernel'): (1, 768, 768)
+('DiTBlock_0', 'Dense_3', 'bias'): (1, 768)
+('DiTBlock_0', 'Dense_3', 'kernel'): (1, 768, 768)
+('DiTBlock_0', 'Dense_4', 'bias'): (1, 768)
+('DiTBlock_0', 'Dense_4', 'kernel'): (1, 768, 768)
+('DiTBlock_0', 'MlpBlock_0', 'Dense_0', 'bias'): (1, 3072)
+('DiTBlock_0', 'MlpBlock_0', 'Dense_0', 'kernel'): (1, 768, 3072)
+('DiTBlock_0', 'MlpBlock_0', 'Dense_1', 'bias'): (1, 768)
+('DiTBlock_0', 'MlpBlock_0', 'Dense_1', 'kernel'): (1, 3072, 768)
+('DiTBlock_1', 'Dense_0', 'bias'): (1, 4608)
+('DiTBlock_1', 'Dense_0', 'kernel'): (1, 768, 4608)
+('DiTBlock_1', 'Dense_1', 'bias'): (1, 768)
+('DiTBlock_1', 'Dense_1', 'kernel'): (1, 768, 768)
+('DiTBlock_1', 'Dense_2', 'bias'): (1, 768)
+('DiTBlock_1', 'Dense_2', 'kernel'): (1, 768, 768)
+('DiTBlock_1', 'Dense_3', 'bias'): (1, 768)
+('DiTBlock_1', 'Dense_3', 'kernel'): (1, 768, 768)
+('DiTBlock_1', 'Dense_4', 'bias'): (1, 768)
+('DiTBlock_1', 'Dense_4', 'kernel'): (1, 768, 768)
+('DiTBlock_1', 'MlpBlock_0', 'Dense_0', 'bias'): (1, 3072)
+('DiTBlock_1', 'MlpBlock_0', 'Dense_0', 'kernel'): (1, 768, 3072)
+('DiTBlock_1', 'MlpBlock_0', 'Dense_1', 'bias'): (1, 768)
+('DiTBlock_1', 'MlpBlock_0', 'Dense_1', 'kernel'): (1, 3072, 768)
+('DiTBlock_10', 'Dense_0', 'bias'): (1, 4608)
+('DiTBlock_10', 'Dense_0', 'kernel'): (1, 768, 4608)
+('DiTBlock_10', 'Dense_1', 'bias'): (1, 768)
+('DiTBlock_10', 'Dense_1', 'kernel'): (1, 768, 768)
+('DiTBlock_10', 'Dense_2', 'bias'): (1, 768)
+('DiTBlock_10', 'Dense_2', 'kernel'): (1, 768, 768)
+('DiTBlock_10', 'Dense_3', 'bias'): (1, 768)
+('DiTBlock_10', 'Dense_3', 'kernel'): (1, 768, 768)
+('DiTBlock_10', 'Dense_4', 'bias'): (1, 768)
+('DiTBlock_10', 'Dense_4', 'kernel'): (1, 768, 768)
+('DiTBlock_10', 'MlpBlock_0', 'Dense_0', 'bias'): (1, 3072)
+('DiTBlock_10', 'MlpBlock_0', 'Dense_0', 'kernel'): (1, 768, 3072)
+('DiTBlock_10', 'MlpBlock_0', 'Dense_1', 'bias'): (1, 768)
+('DiTBlock_10', 'MlpBlock_0', 'Dense_1', 'kernel'): (1, 3072, 768)
+('DiTBlock_11', 'Dense_0', 'bias'): (1, 4608)
+('DiTBlock_11', 'Dense_0', 'kernel'): (1, 768, 4608)
+('DiTBlock_11', 'Dense_1', 'bias'): (1, 768)
+('DiTBlock_11', 'Dense_1', 'kernel'): (1, 768, 768)
+('DiTBlock_11', 'Dense_2', 'bias'): (1, 768)
+('DiTBlock_11', 'Dense_2', 'kernel'): (1, 768, 768)
+('DiTBlock_11', 'Dense_3', 'bias'): (1, 768)
+('DiTBlock_11', 'Dense_3', 'kernel'): (1, 768, 768)
+('DiTBlock_11', 'Dense_4', 'bias'): (1, 768)
+('DiTBlock_11', 'Dense_4', 'kernel'): (1, 768, 768)
+('DiTBlock_11', 'MlpBlock_0', 'Dense_0', 'bias'): (1, 3072)
+('DiTBlock_11', 'MlpBlock_0', 'Dense_0', 'kernel'): (1, 768, 3072)
+('DiTBlock_11', 'MlpBlock_0', 'Dense_1', 'bias'): (1, 768)
+('DiTBlock_11', 'MlpBlock_0', 'Dense_1', 'kernel'): (1, 3072, 768)
+('DiTBlock_2', 'Dense_0', 'bias'): (1, 4608)
+('DiTBlock_2', 'Dense_0', 'kernel'): (1, 768, 4608)
+('DiTBlock_2', 'Dense_1', 'bias'): (1, 768)
+('DiTBlock_2', 'Dense_1', 'kernel'): (1, 768, 768)
+('DiTBlock_2', 'Dense_2', 'bias'): (1, 768)
+('DiTBlock_2', 'Dense_2', 'kernel'): (1, 768, 768)
+('DiTBlock_2', 'Dense_3', 'bias'): (1, 768)
+('DiTBlock_2', 'Dense_3', 'kernel'): (1, 768, 768)
+('DiTBlock_2', 'Dense_4', 'bias'): (1, 768)
+('DiTBlock_2', 'Dense_4', 'kernel'): (1, 768, 768)
+('DiTBlock_2', 'MlpBlock_0', 'Dense_0', 'bias'): (1, 3072)
+('DiTBlock_2', 'MlpBlock_0', 'Dense_0', 'kernel'): (1, 768, 3072)
+('DiTBlock_2', 'MlpBlock_0', 'Dense_1', 'bias'): (1, 768)
+('DiTBlock_2', 'MlpBlock_0', 'Dense_1', 'kernel'): (1, 3072, 768)
+('DiTBlock_3', 'Dense_0', 'bias'): (1, 4608)
+('DiTBlock_3', 'Dense_0', 'kernel'): (1, 768, 4608)
+('DiTBlock_3', 'Dense_1', 'bias'): (1, 768)
+('DiTBlock_3', 'Dense_1', 'kernel'): (1, 768, 768)
+('DiTBlock_3', 'Dense_2', 'bias'): (1, 768)
+('DiTBlock_3', 'Dense_2', 'kernel'): (1, 768, 768)
+('DiTBlock_3', 'Dense_3', 'bias'): (1, 768)
+('DiTBlock_3', 'Dense_3', 'kernel'): (1, 768, 768)
+('DiTBlock_3', 'Dense_4', 'bias'): (1, 768)
+('DiTBlock_3', 'Dense_4', 'kernel'): (1, 768, 768)
+('DiTBlock_3', 'MlpBlock_0', 'Dense_0', 'bias'): (1, 3072)
+('DiTBlock_3', 'MlpBlock_0', 'Dense_0', 'kernel'): (1, 768, 3072)
+('DiTBlock_3', 'MlpBlock_0', 'Dense_1', 'bias'): (1, 768)
+('DiTBlock_3', 'MlpBlock_0', 'Dense_1', 'kernel'): (1, 3072, 768)
+('DiTBlock_4', 'Dense_0', 'bias'): (1, 4608)
+('DiTBlock_4', 'Dense_0', 'kernel'): (1, 768, 4608)
+('DiTBlock_4', 'Dense_1', 'bias'): (1, 768)
+('DiTBlock_4', 'Dense_1', 'kernel'): (1, 768, 768)
+('DiTBlock_4', 'Dense_2', 'bias'): (1, 768)
+('DiTBlock_4', 'Dense_2', 'kernel'): (1, 768, 768)
+('DiTBlock_4', 'Dense_3', 'bias'): (1, 768)
+('DiTBlock_4', 'Dense_3', 'kernel'): (1, 768, 768)
+('DiTBlock_4', 'Dense_4', 'bias'): (1, 768)
+('DiTBlock_4', 'Dense_4', 'kernel'): (1, 768, 768)
+('DiTBlock_4', 'MlpBlock_0', 'Dense_0', 'bias'): (1, 3072)
+('DiTBlock_4', 'MlpBlock_0', 'Dense_0', 'kernel'): (1, 768, 3072)
+('DiTBlock_4', 'MlpBlock_0', 'Dense_1', 'bias'): (1, 768)
+('DiTBlock_4', 'MlpBlock_0', 'Dense_1', 'kernel'): (1, 3072, 768)
+('DiTBlock_5', 'Dense_0', 'bias'): (1, 4608)
+('DiTBlock_5', 'Dense_0', 'kernel'): (1, 768, 4608)
+('DiTBlock_5', 'Dense_1', 'bias'): (1, 768)
+('DiTBlock_5', 'Dense_1', 'kernel'): (1, 768, 768)
+('DiTBlock_5', 'Dense_2', 'bias'): (1, 768)
+('DiTBlock_5', 'Dense_2', 'kernel'): (1, 768, 768)
+('DiTBlock_5', 'Dense_3', 'bias'): (1, 768)
+('DiTBlock_5', 'Dense_3', 'kernel'): (1, 768, 768)
+('DiTBlock_5', 'Dense_4', 'bias'): (1, 768)
+('DiTBlock_5', 'Dense_4', 'kernel'): (1, 768, 768)
+('DiTBlock_5', 'MlpBlock_0', 'Dense_0', 'bias'): (1, 3072)
+('DiTBlock_5', 'MlpBlock_0', 'Dense_0', 'kernel'): (1, 768, 3072)
+('DiTBlock_5', 'MlpBlock_0', 'Dense_1', 'bias'): (1, 768)
+('DiTBlock_5', 'MlpBlock_0', 'Dense_1', 'kernel'): (1, 3072, 768)
+('DiTBlock_6', 'Dense_0', 'bias'): (1, 4608)
+('DiTBlock_6', 'Dense_0', 'kernel'): (1, 768, 4608)
+('DiTBlock_6', 'Dense_1', 'bias'): (1, 768)
+('DiTBlock_6', 'Dense_1', 'kernel'): (1, 768, 768)
+('DiTBlock_6', 'Dense_2', 'bias'): (1, 768)
+('DiTBlock_6', 'Dense_2', 'kernel'): (1, 768, 768)
+('DiTBlock_6', 'Dense_3', 'bias'): (1, 768)
+('DiTBlock_6', 'Dense_3', 'kernel'): (1, 768, 768)
+('DiTBlock_6', 'Dense_4', 'bias'): (1, 768)
+('DiTBlock_6', 'Dense_4', 'kernel'): (1, 768, 768)
+('DiTBlock_6', 'MlpBlock_0', 'Dense_0', 'bias'): (1, 3072)
+('DiTBlock_6', 'MlpBlock_0', 'Dense_0', 'kernel'): (1, 768, 3072)
+('DiTBlock_6', 'MlpBlock_0', 'Dense_1', 'bias'): (1, 768)
+('DiTBlock_6', 'MlpBlock_0', 'Dense_1', 'kernel'): (1, 3072, 768)
+('DiTBlock_7', 'Dense_0', 'bias'): (1, 4608)
+('DiTBlock_7', 'Dense_0', 'kernel'): (1, 768, 4608)
+('DiTBlock_7', 'Dense_1', 'bias'): (1, 768)
+('DiTBlock_7', 'Dense_1', 'kernel'): (1, 768, 768)
+('DiTBlock_7', 'Dense_2', 'bias'): (1, 768)
+('DiTBlock_7', 'Dense_2', 'kernel'): (1, 768, 768)
+('DiTBlock_7', 'Dense_3', 'bias'): (1, 768)
+('DiTBlock_7', 'Dense_3', 'kernel'): (1, 768, 768)
+('DiTBlock_7', 'Dense_4', 'bias'): (1, 768)
+('DiTBlock_7', 'Dense_4', 'kernel'): (1, 768, 768)
+('DiTBlock_7', 'MlpBlock_0', 'Dense_0', 'bias'): (1, 3072)
+('DiTBlock_7', 'MlpBlock_0', 'Dense_0', 'kernel'): (1, 768, 3072)
+('DiTBlock_7', 'MlpBlock_0', 'Dense_1', 'bias'): (1, 768)
+('DiTBlock_7', 'MlpBlock_0', 'Dense_1', 'kernel'): (1, 3072, 768)
+('DiTBlock_8', 'Dense_0', 'bias'): (1, 4608)
+('DiTBlock_8', 'Dense_0', 'kernel'): (1, 768, 4608)
+('DiTBlock_8', 'Dense_1', 'bias'): (1, 768)
+('DiTBlock_8', 'Dense_1', 'kernel'): (1, 768, 768)
+('DiTBlock_8', 'Dense_2', 'bias'): (1, 768)
+('DiTBlock_8', 'Dense_2', 'kernel'): (1, 768, 768)
+('DiTBlock_8', 'Dense_3', 'bias'): (1, 768)
+('DiTBlock_8', 'Dense_3', 'kernel'): (1, 768, 768)
+('DiTBlock_8', 'Dense_4', 'bias'): (1, 768)
+('DiTBlock_8', 'Dense_4', 'kernel'): (1, 768, 768)
+('DiTBlock_8', 'MlpBlock_0', 'Dense_0', 'bias'): (1, 3072)
+('DiTBlock_8', 'MlpBlock_0', 'Dense_0', 'kernel'): (1, 768, 3072)
+('DiTBlock_8', 'MlpBlock_0', 'Dense_1', 'bias'): (1, 768)
+('DiTBlock_8', 'MlpBlock_0', 'Dense_1', 'kernel'): (1, 3072, 768)
+('DiTBlock_9', 'Dense_0', 'bias'): (1, 4608)
+('DiTBlock_9', 'Dense_0', 'kernel'): (1, 768, 4608)
+('DiTBlock_9', 'Dense_1', 'bias'): (1, 768)
+('DiTBlock_9', 'Dense_1', 'kernel'): (1, 768, 768)
+('DiTBlock_9', 'Dense_2', 'bias'): (1, 768)
+('DiTBlock_9', 'Dense_2', 'kernel'): (1, 768, 768)
+('DiTBlock_9', 'Dense_3', 'bias'): (1, 768)
+('DiTBlock_9', 'Dense_3', 'kernel'): (1, 768, 768)
+('DiTBlock_9', 'Dense_4', 'bias'): (1, 768)
+('DiTBlock_9', 'Dense_4', 'kernel'): (1, 768, 768)
+('DiTBlock_9', 'MlpBlock_0', 'Dense_0', 'bias'): (1, 3072)
+('DiTBlock_9', 'MlpBlock_0', 'Dense_0', 'kernel'): (1, 768, 3072)
+('DiTBlock_9', 'MlpBlock_0', 'Dense_1', 'bias'): (1, 768)
+('DiTBlock_9', 'MlpBlock_0', 'Dense_1', 'kernel'): (1, 3072, 768)
+('Embed_0', 'embedding'): (1, 256, 1)
+('FinalLayer_0', 'Dense_0', 'bias'): (1, 1536)
+('FinalLayer_0', 'Dense_0', 'kernel'): (1, 768, 1536)
+('FinalLayer_0', 'Dense_1', 'bias'): (1, 16)
+('FinalLayer_0', 'Dense_1', 'kernel'): (1, 768, 16)
+('LabelEmbedder_0', 'Embed_0', 'embedding'): (1, 1001, 768)
+('PatchEmbed_0', 'Conv_0', 'bias'): (1, 768)
+('PatchEmbed_0', 'Conv_0', 'kernel'): (1, 2, 2, 4, 768)
+('TimestepEmbedder_0', 'Dense_0', 'bias'): (1, 768)
+('TimestepEmbedder_0', 'Dense_0', 'kernel'): (1, 256, 768)
+('TimestepEmbedder_0', 'Dense_1', 'bias'): (1, 768)
+('TimestepEmbedder_0', 'Dense_1', 'kernel'): (1, 768, 768)
+('TimestepEmbedder_1', 'Dense_0', 'bias'): (1, 768)
+('TimestepEmbedder_1', 'Dense_0', 'kernel'): (1, 256, 768)
+('TimestepEmbedder_1', 'Dense_1', 'bias'): (1, 768)
+('TimestepEmbedder_1', 'Dense_1', 'kernel'): (1, 768, 768)
+
+ parameter shapes:
+('DiTBlock_0', 'Dense_0', 'bias'): (1, 4608)
+('DiTBlock_0', 'Dense_0', 'kernel'): (1, 768, 4608)
+('DiTBlock_0', 'Dense_1', 'bias'): (1, 768)
+('DiTBlock_0', 'Dense_1', 'kernel'): (1, 768, 768)
+('DiTBlock_0', 'Dense_2', 'bias'): (1, 768)
+('DiTBlock_0', 'Dense_2', 'kernel'): (1, 768, 768)
+('DiTBlock_0', 'Dense_3', 'bias'): (1, 768)
+('DiTBlock_0', 'Dense_3', 'kernel'): (1, 768, 768)
+('DiTBlock_0', 'Dense_4', 'bias'): (1, 768)
+('DiTBlock_0', 'Dense_4', 'kernel'): (1, 768, 768)
+('DiTBlock_0', 'MlpBlock_0', 'Dense_0', 'bias'): (1, 3072)
+('DiTBlock_0', 'MlpBlock_0', 'Dense_0', 'kernel'): (1, 768, 3072)
+('DiTBlock_0', 'MlpBlock_0', 'Dense_1', 'bias'): (1, 768)
+('DiTBlock_0', 'MlpBlock_0', 'Dense_1', 'kernel'): (1, 3072, 768)
+('DiTBlock_1', 'Dense_0', 'bias'): (1, 4608)
+('DiTBlock_1', 'Dense_0', 'kernel'): (1, 768, 4608)
+('DiTBlock_1', 'Dense_1', 'bias'): (1, 768)
+('DiTBlock_1', 'Dense_1', 'kernel'): (1, 768, 768)
+('DiTBlock_1', 'Dense_2', 'bias'): (1, 768)
+('DiTBlock_1', 'Dense_2', 'kernel'): (1, 768, 768)
+('DiTBlock_1', 'Dense_3', 'bias'): (1, 768)
+('DiTBlock_1', 'Dense_3', 'kernel'): (1, 768, 768)
+('DiTBlock_1', 'Dense_4', 'bias'): (1, 768)
+('DiTBlock_1', 'Dense_4', 'kernel'): (1, 768, 768)
+('DiTBlock_1', 'MlpBlock_0', 'Dense_0', 'bias'): (1, 3072)
+('DiTBlock_1', 'MlpBlock_0', 'Dense_0', 'kernel'): (1, 768, 3072)
+('DiTBlock_1', 'MlpBlock_0', 'Dense_1', 'bias'): (1, 768)
+('DiTBlock_1', 'MlpBlock_0', 'Dense_1', 'kernel'): (1, 3072, 768)
+('DiTBlock_10', 'Dense_0', 'bias'): (1, 4608)
+('DiTBlock_10', 'Dense_0', 'kernel'): (1, 768, 4608)
+('DiTBlock_10', 'Dense_1', 'bias'): (1, 768)
+('DiTBlock_10', 'Dense_1', 'kernel'): (1, 768, 768)
+('DiTBlock_10', 'Dense_2', 'bias'): (1, 768)
+('DiTBlock_10', 'Dense_2', 'kernel'): (1, 768, 768)
+('DiTBlock_10', 'Dense_3', 'bias'): (1, 768)
+('DiTBlock_10', 'Dense_3', 'kernel'): (1, 768, 768)
+('DiTBlock_10', 'Dense_4', 'bias'): (1, 768)
+('DiTBlock_10', 'Dense_4', 'kernel'): (1, 768, 768)
+('DiTBlock_10', 'MlpBlock_0', 'Dense_0', 'bias'): (1, 3072)
+('DiTBlock_10', 'MlpBlock_0', 'Dense_0', 'kernel'): (1, 768, 3072)
+('DiTBlock_10', 'MlpBlock_0', 'Dense_1', 'bias'): (1, 768)
+('DiTBlock_10', 'MlpBlock_0', 'Dense_1', 'kernel'): (1, 3072, 768)
+('DiTBlock_11', 'Dense_0', 'bias'): (1, 4608)
+('DiTBlock_11', 'Dense_0', 'kernel'): (1, 768, 4608)
+('DiTBlock_11', 'Dense_1', 'bias'): (1, 768)
+('DiTBlock_11', 'Dense_1', 'kernel'): (1, 768, 768)
+('DiTBlock_11', 'Dense_2', 'bias'): (1, 768)
+('DiTBlock_11', 'Dense_2', 'kernel'): (1, 768, 768)
+('DiTBlock_11', 'Dense_3', 'bias'): (1, 768)
+('DiTBlock_11', 'Dense_3', 'kernel'): (1, 768, 768)
+('DiTBlock_11', 'Dense_4', 'bias'): (1, 768)
+('DiTBlock_11', 'Dense_4', 'kernel'): (1, 768, 768)
+('DiTBlock_11', 'MlpBlock_0', 'Dense_0', 'bias'): (1, 3072)
+('DiTBlock_11', 'MlpBlock_0', 'Dense_0', 'kernel'): (1, 768, 3072)
+('DiTBlock_11', 'MlpBlock_0', 'Dense_1', 'bias'): (1, 768)
+('DiTBlock_11', 'MlpBlock_0', 'Dense_1', 'kernel'): (1, 3072, 768)
+('DiTBlock_2', 'Dense_0', 'bias'): (1, 4608)
+('DiTBlock_2', 'Dense_0', 'kernel'): (1, 768, 4608)
+('DiTBlock_2', 'Dense_1', 'bias'): (1, 768)
+('DiTBlock_2', 'Dense_1', 'kernel'): (1, 768, 768)
+('DiTBlock_2', 'Dense_2', 'bias'): (1, 768)
+('DiTBlock_2', 'Dense_2', 'kernel'): (1, 768, 768)
+('DiTBlock_2', 'Dense_3', 'bias'): (1, 768)
+('DiTBlock_2', 'Dense_3', 'kernel'): (1, 768, 768)
+('DiTBlock_2', 'Dense_4', 'bias'): (1, 768)
+('DiTBlock_2', 'Dense_4', 'kernel'): (1, 768, 768)
+('DiTBlock_2', 'MlpBlock_0', 'Dense_0', 'bias'): (1, 3072)
+('DiTBlock_2', 'MlpBlock_0', 'Dense_0', 'kernel'): (1, 768, 3072)
+('DiTBlock_2', 'MlpBlock_0', 'Dense_1', 'bias'): (1, 768)
+('DiTBlock_2', 'MlpBlock_0', 'Dense_1', 'kernel'): (1, 3072, 768)
+('DiTBlock_3', 'Dense_0', 'bias'): (1, 4608)
+('DiTBlock_3', 'Dense_0', 'kernel'): (1, 768, 4608)
+('DiTBlock_3', 'Dense_1', 'bias'): (1, 768)
+('DiTBlock_3', 'Dense_1', 'kernel'): (1, 768, 768)
+('DiTBlock_3', 'Dense_2', 'bias'): (1, 768)
+('DiTBlock_3', 'Dense_2', 'kernel'): (1, 768, 768)
+('DiTBlock_3', 'Dense_3', 'bias'): (1, 768)
+('DiTBlock_3', 'Dense_3', 'kernel'): (1, 768, 768)
+('DiTBlock_3', 'Dense_4', 'bias'): (1, 768)
+('DiTBlock_3', 'Dense_4', 'kernel'): (1, 768, 768)
+('DiTBlock_3', 'MlpBlock_0', 'Dense_0', 'bias'): (1, 3072)
+('DiTBlock_3', 'MlpBlock_0', 'Dense_0', 'kernel'): (1, 768, 3072)
+('DiTBlock_3', 'MlpBlock_0', 'Dense_1', 'bias'): (1, 768)
+('DiTBlock_3', 'MlpBlock_0', 'Dense_1', 'kernel'): (1, 3072, 768)
+('DiTBlock_4', 'Dense_0', 'bias'): (1, 4608)
+('DiTBlock_4', 'Dense_0', 'kernel'): (1, 768, 4608)
+('DiTBlock_4', 'Dense_1', 'bias'): (1, 768)
+('DiTBlock_4', 'Dense_1', 'kernel'): (1, 768, 768)
+('DiTBlock_4', 'Dense_2', 'bias'): (1, 768)
+('DiTBlock_4', 'Dense_2', 'kernel'): (1, 768, 768)
+('DiTBlock_4', 'Dense_3', 'bias'): (1, 768)
+('DiTBlock_4', 'Dense_3', 'kernel'): (1, 768, 768)
+('DiTBlock_4', 'Dense_4', 'bias'): (1, 768)
+('DiTBlock_4', 'Dense_4', 'kernel'): (1, 768, 768)
+('DiTBlock_4', 'MlpBlock_0', 'Dense_0', 'bias'): (1, 3072)
+('DiTBlock_4', 'MlpBlock_0', 'Dense_0', 'kernel'): (1, 768, 3072)
+('DiTBlock_4', 'MlpBlock_0', 'Dense_1', 'bias'): (1, 768)
+('DiTBlock_4', 'MlpBlock_0', 'Dense_1', 'kernel'): (1, 3072, 768)
+('DiTBlock_5', 'Dense_0', 'bias'): (1, 4608)
+('DiTBlock_5', 'Dense_0', 'kernel'): (1, 768, 4608)
+('DiTBlock_5', 'Dense_1', 'bias'): (1, 768)
+('DiTBlock_5', 'Dense_1', 'kernel'): (1, 768, 768)
+('DiTBlock_5', 'Dense_2', 'bias'): (1, 768)
+('DiTBlock_5', 'Dense_2', 'kernel'): (1, 768, 768)
+('DiTBlock_5', 'Dense_3', 'bias'): (1, 768)
+('DiTBlock_5', 'Dense_3', 'kernel'): (1, 768, 768)
+('DiTBlock_5', 'Dense_4', 'bias'): (1, 768)
+('DiTBlock_5', 'Dense_4', 'kernel'): (1, 768, 768)
+('DiTBlock_5', 'MlpBlock_0', 'Dense_0', 'bias'): (1, 3072)
+('DiTBlock_5', 'MlpBlock_0', 'Dense_0', 'kernel'): (1, 768, 3072)
+('DiTBlock_5', 'MlpBlock_0', 'Dense_1', 'bias'): (1, 768)
+('DiTBlock_5', 'MlpBlock_0', 'Dense_1', 'kernel'): (1, 3072, 768)
+('DiTBlock_6', 'Dense_0', 'bias'): (1, 4608)
+('DiTBlock_6', 'Dense_0', 'kernel'): (1, 768, 4608)
+('DiTBlock_6', 'Dense_1', 'bias'): (1, 768)
+('DiTBlock_6', 'Dense_1', 'kernel'): (1, 768, 768)
+('DiTBlock_6', 'Dense_2', 'bias'): (1, 768)
+('DiTBlock_6', 'Dense_2', 'kernel'): (1, 768, 768)
+('DiTBlock_6', 'Dense_3', 'bias'): (1, 768)
+('DiTBlock_6', 'Dense_3', 'kernel'): (1, 768, 768)
+('DiTBlock_6', 'Dense_4', 'bias'): (1, 768)
+('DiTBlock_6', 'Dense_4', 'kernel'): (1, 768, 768)
+('DiTBlock_6', 'MlpBlock_0', 'Dense_0', 'bias'): (1, 3072)
+('DiTBlock_6', 'MlpBlock_0', 'Dense_0', 'kernel'): (1, 768, 3072)
+('DiTBlock_6', 'MlpBlock_0', 'Dense_1', 'bias'): (1, 768)
+('DiTBlock_6', 'MlpBlock_0', 'Dense_1', 'kernel'): (1, 3072, 768)
+('DiTBlock_7', 'Dense_0', 'bias'): (1, 4608)
+('DiTBlock_7', 'Dense_0', 'kernel'): (1, 768, 4608)
+('DiTBlock_7', 'Dense_1', 'bias'): (1, 768)
+('DiTBlock_7', 'Dense_1', 'kernel'): (1, 768, 768)
+('DiTBlock_7', 'Dense_2', 'bias'): (1, 768)
+('DiTBlock_7', 'Dense_2', 'kernel'): (1, 768, 768)
+('DiTBlock_7', 'Dense_3', 'bias'): (1, 768)
+('DiTBlock_7', 'Dense_3', 'kernel'): (1, 768, 768)
+('DiTBlock_7', 'Dense_4', 'bias'): (1, 768)
+('DiTBlock_7', 'Dense_4', 'kernel'): (1, 768, 768)
+('DiTBlock_7', 'MlpBlock_0', 'Dense_0', 'bias'): (1, 3072)
+('DiTBlock_7', 'MlpBlock_0', 'Dense_0', 'kernel'): (1, 768, 3072)
+('DiTBlock_7', 'MlpBlock_0', 'Dense_1', 'bias'): (1, 768)
+('DiTBlock_7', 'MlpBlock_0', 'Dense_1', 'kernel'): (1, 3072, 768)
+('DiTBlock_8', 'Dense_0', 'bias'): (1, 4608)
+('DiTBlock_8', 'Dense_0', 'kernel'): (1, 768, 4608)
+('DiTBlock_8', 'Dense_1', 'bias'): (1, 768)
+('DiTBlock_8', 'Dense_1', 'kernel'): (1, 768, 768)
+('DiTBlock_8', 'Dense_2', 'bias'): (1, 768)
+('DiTBlock_8', 'Dense_2', 'kernel'): (1, 768, 768)
+('DiTBlock_8', 'Dense_3', 'bias'): (1, 768)
+('DiTBlock_8', 'Dense_3', 'kernel'): (1, 768, 768)
+('DiTBlock_8', 'Dense_4', 'bias'): (1, 768)
+('DiTBlock_8', 'Dense_4', 'kernel'): (1, 768, 768)
+('DiTBlock_8', 'MlpBlock_0', 'Dense_0', 'bias'): (1, 3072)
+('DiTBlock_8', 'MlpBlock_0', 'Dense_0', 'kernel'): (1, 768, 3072)
+('DiTBlock_8', 'MlpBlock_0', 'Dense_1', 'bias'): (1, 768)
+('DiTBlock_8', 'MlpBlock_0', 'Dense_1', 'kernel'): (1, 3072, 768)
+('DiTBlock_9', 'Dense_0', 'bias'): (1, 4608)
+('DiTBlock_9', 'Dense_0', 'kernel'): (1, 768, 4608)
+('DiTBlock_9', 'Dense_1', 'bias'): (1, 768)
+('DiTBlock_9', 'Dense_1', 'kernel'): (1, 768, 768)
+('DiTBlock_9', 'Dense_2', 'bias'): (1, 768)
+('DiTBlock_9', 'Dense_2', 'kernel'): (1, 768, 768)
+('DiTBlock_9', 'Dense_3', 'bias'): (1, 768)
+('DiTBlock_9', 'Dense_3', 'kernel'): (1, 768, 768)
+('DiTBlock_9', 'Dense_4', 'bias'): (1, 768)
+('DiTBlock_9', 'Dense_4', 'kernel'): (1, 768, 768)
+('DiTBlock_9', 'MlpBlock_0', 'Dense_0', 'bias'): (1, 3072)
+('DiTBlock_9', 'MlpBlock_0', 'Dense_0', 'kernel'): (1, 768, 3072)
+('DiTBlock_9', 'MlpBlock_0', 'Dense_1', 'bias'): (1, 768)
+('DiTBlock_9', 'MlpBlock_0', 'Dense_1', 'kernel'): (1, 3072, 768)
+('Embed_0', 'embedding'): (1, 256, 1)
+('FinalLayer_0', 'Dense_0', 'bias'): (1, 1536)
+('FinalLayer_0', 'Dense_0', 'kernel'): (1, 768, 1536)
+('FinalLayer_0', 'Dense_1', 'bias'): (1, 16)
+('FinalLayer_0', 'Dense_1', 'kernel'): (1, 768, 16)
+('LabelEmbedder_0', 'Embed_0', 'embedding'): (1, 1001, 768)
+('PatchEmbed_0', 'Conv_0', 'bias'): (1, 768)
+('PatchEmbed_0', 'Conv_0', 'kernel'): (1, 2, 2, 4, 768)
+('TimestepEmbedder_0', 'Dense_0', 'bias'): (1, 768)
+('TimestepEmbedder_0', 'Dense_0', 'kernel'): (1, 256, 768)
+('TimestepEmbedder_0', 'Dense_1', 'bias'): (1, 768)
+('TimestepEmbedder_0', 'Dense_1', 'kernel'): (1, 768, 768)
+('TimestepEmbedder_1', 'Dense_0', 'bias'): (1, 768)
+('TimestepEmbedder_1', 'Dense_0', 'kernel'): (1, 256, 768)
+('TimestepEmbedder_1', 'Dense_1', 'bias'): (1, 768)
+('TimestepEmbedder_1', 'Dense_1', 'kernel'): (1, 768, 768)
+
+ parameter shapes:
+('DiTBlock_0', 'Dense_0', 'bias'): (1, 4608)
+('DiTBlock_0', 'Dense_0', 'kernel'): (1, 768, 4608)
+('DiTBlock_0', 'Dense_1', 'bias'): (1, 768)
+('DiTBlock_0', 'Dense_1', 'kernel'): (1, 768, 768)
+('DiTBlock_0', 'Dense_2', 'bias'): (1, 768)
+('DiTBlock_0', 'Dense_2', 'kernel'): (1, 768, 768)
+('DiTBlock_0', 'Dense_3', 'bias'): (1, 768)
+('DiTBlock_0', 'Dense_3', 'kernel'): (1, 768, 768)
+('DiTBlock_0', 'Dense_4', 'bias'): (1, 768)
+('DiTBlock_0', 'Dense_4', 'kernel'): (1, 768, 768)
+('DiTBlock_0', 'MlpBlock_0', 'Dense_0', 'bias'): (1, 3072)
+('DiTBlock_0', 'MlpBlock_0', 'Dense_0', 'kernel'): (1, 768, 3072)
+('DiTBlock_0', 'MlpBlock_0', 'Dense_1', 'bias'): (1, 768)
+('DiTBlock_0', 'MlpBlock_0', 'Dense_1', 'kernel'): (1, 3072, 768)
+('DiTBlock_1', 'Dense_0', 'bias'): (1, 4608)
+('DiTBlock_1', 'Dense_0', 'kernel'): (1, 768, 4608)
+('DiTBlock_1', 'Dense_1', 'bias'): (1, 768)
+('DiTBlock_1', 'Dense_1', 'kernel'): (1, 768, 768)
+('DiTBlock_1', 'Dense_2', 'bias'): (1, 768)
+('DiTBlock_1', 'Dense_2', 'kernel'): (1, 768, 768)
+('DiTBlock_1', 'Dense_3', 'bias'): (1, 768)
+('DiTBlock_1', 'Dense_3', 'kernel'): (1, 768, 768)
+('DiTBlock_1', 'Dense_4', 'bias'): (1, 768)
+('DiTBlock_1', 'Dense_4', 'kernel'): (1, 768, 768)
+('DiTBlock_1', 'MlpBlock_0', 'Dense_0', 'bias'): (1, 3072)
+('DiTBlock_1', 'MlpBlock_0', 'Dense_0', 'kernel'): (1, 768, 3072)
+('DiTBlock_1', 'MlpBlock_0', 'Dense_1', 'bias'): (1, 768)
+('DiTBlock_1', 'MlpBlock_0', 'Dense_1', 'kernel'): (1, 3072, 768)
+('DiTBlock_10', 'Dense_0', 'bias'): (1, 4608)
+('DiTBlock_10', 'Dense_0', 'kernel'): (1, 768, 4608)
+('DiTBlock_10', 'Dense_1', 'bias'): (1, 768)
+('DiTBlock_10', 'Dense_1', 'kernel'): (1, 768, 768)
+('DiTBlock_10', 'Dense_2', 'bias'): (1, 768)
+('DiTBlock_10', 'Dense_2', 'kernel'): (1, 768, 768)
+('DiTBlock_10', 'Dense_3', 'bias'): (1, 768)
+('DiTBlock_10', 'Dense_3', 'kernel'): (1, 768, 768)
+('DiTBlock_10', 'Dense_4', 'bias'): (1, 768)
+('DiTBlock_10', 'Dense_4', 'kernel'): (1, 768, 768)
+('DiTBlock_10', 'MlpBlock_0', 'Dense_0', 'bias'): (1, 3072)
+('DiTBlock_10', 'MlpBlock_0', 'Dense_0', 'kernel'): (1, 768, 3072)
+('DiTBlock_10', 'MlpBlock_0', 'Dense_1', 'bias'): (1, 768)
+('DiTBlock_10', 'MlpBlock_0', 'Dense_1', 'kernel'): (1, 3072, 768)
+('DiTBlock_11', 'Dense_0', 'bias'): (1, 4608)
+('DiTBlock_11', 'Dense_0', 'kernel'): (1, 768, 4608)
+('DiTBlock_11', 'Dense_1', 'bias'): (1, 768)
+('DiTBlock_11', 'Dense_1', 'kernel'): (1, 768, 768)
+('DiTBlock_11', 'Dense_2', 'bias'): (1, 768)
+('DiTBlock_11', 'Dense_2', 'kernel'): (1, 768, 768)
+('DiTBlock_11', 'Dense_3', 'bias'): (1, 768)
+('DiTBlock_11', 'Dense_3', 'kernel'): (1, 768, 768)
+('DiTBlock_11', 'Dense_4', 'bias'): (1, 768)
+('DiTBlock_11', 'Dense_4', 'kernel'): (1, 768, 768)
+('DiTBlock_11', 'MlpBlock_0', 'Dense_0', 'bias'): (1, 3072)
+('DiTBlock_11', 'MlpBlock_0', 'Dense_0', 'kernel'): (1, 768, 3072)
+('DiTBlock_11', 'MlpBlock_0', 'Dense_1', 'bias'): (1, 768)
+('DiTBlock_11', 'MlpBlock_0', 'Dense_1', 'kernel'): (1, 3072, 768)
+('DiTBlock_2', 'Dense_0', 'bias'): (1, 4608)
+('DiTBlock_2', 'Dense_0', 'kernel'): (1, 768, 4608)
+('DiTBlock_2', 'Dense_1', 'bias'): (1, 768)
+('DiTBlock_2', 'Dense_1', 'kernel'): (1, 768, 768)
+('DiTBlock_2', 'Dense_2', 'bias'): (1, 768)
+('DiTBlock_2', 'Dense_2', 'kernel'): (1, 768, 768)
+('DiTBlock_2', 'Dense_3', 'bias'): (1, 768)
+('DiTBlock_2', 'Dense_3', 'kernel'): (1, 768, 768)
+('DiTBlock_2', 'Dense_4', 'bias'): (1, 768)
+('DiTBlock_2', 'Dense_4', 'kernel'): (1, 768, 768)
+('DiTBlock_2', 'MlpBlock_0', 'Dense_0', 'bias'): (1, 3072)
+('DiTBlock_2', 'MlpBlock_0', 'Dense_0', 'kernel'): (1, 768, 3072)
+('DiTBlock_2', 'MlpBlock_0', 'Dense_1', 'bias'): (1, 768)
+('DiTBlock_2', 'MlpBlock_0', 'Dense_1', 'kernel'): (1, 3072, 768)
+('DiTBlock_3', 'Dense_0', 'bias'): (1, 4608)
+('DiTBlock_3', 'Dense_0', 'kernel'): (1, 768, 4608)
+('DiTBlock_3', 'Dense_1', 'bias'): (1, 768)
+('DiTBlock_3', 'Dense_1', 'kernel'): (1, 768, 768)
+('DiTBlock_3', 'Dense_2', 'bias'): (1, 768)
+('DiTBlock_3', 'Dense_2', 'kernel'): (1, 768, 768)
+('DiTBlock_3', 'Dense_3', 'bias'): (1, 768)
+('DiTBlock_3', 'Dense_3', 'kernel'): (1, 768, 768)
+('DiTBlock_3', 'Dense_4', 'bias'): (1, 768)
+('DiTBlock_3', 'Dense_4', 'kernel'): (1, 768, 768)
+('DiTBlock_3', 'MlpBlock_0', 'Dense_0', 'bias'): (1, 3072)
+('DiTBlock_3', 'MlpBlock_0', 'Dense_0', 'kernel'): (1, 768, 3072)
+('DiTBlock_3', 'MlpBlock_0', 'Dense_1', 'bias'): (1, 768)
+('DiTBlock_3', 'MlpBlock_0', 'Dense_1', 'kernel'): (1, 3072, 768)
+('DiTBlock_4', 'Dense_0', 'bias'): (1, 4608)
+('DiTBlock_4', 'Dense_0', 'kernel'): (1, 768, 4608)
+('DiTBlock_4', 'Dense_1', 'bias'): (1, 768)
+('DiTBlock_4', 'Dense_1', 'kernel'): (1, 768, 768)
+('DiTBlock_4', 'Dense_2', 'bias'): (1, 768)
+('DiTBlock_4', 'Dense_2', 'kernel'): (1, 768, 768)
+('DiTBlock_4', 'Dense_3', 'bias'): (1, 768)
+('DiTBlock_4', 'Dense_3', 'kernel'): (1, 768, 768)
+('DiTBlock_4', 'Dense_4', 'bias'): (1, 768)
+('DiTBlock_4', 'Dense_4', 'kernel'): (1, 768, 768)
+('DiTBlock_4', 'MlpBlock_0', 'Dense_0', 'bias'): (1, 3072)
+('DiTBlock_4', 'MlpBlock_0', 'Dense_0', 'kernel'): (1, 768, 3072)
+('DiTBlock_4', 'MlpBlock_0', 'Dense_1', 'bias'): (1, 768)
+('DiTBlock_4', 'MlpBlock_0', 'Dense_1', 'kernel'): (1, 3072, 768)
+('DiTBlock_5', 'Dense_0', 'bias'): (1, 4608)
+('DiTBlock_5', 'Dense_0', 'kernel'): (1, 768, 4608)
+('DiTBlock_5', 'Dense_1', 'bias'): (1, 768)
+('DiTBlock_5', 'Dense_1', 'kernel'): (1, 768, 768)
+('DiTBlock_5', 'Dense_2', 'bias'): (1, 768)
+('DiTBlock_5', 'Dense_2', 'kernel'): (1, 768, 768)
+('DiTBlock_5', 'Dense_3', 'bias'): (1, 768)
+('DiTBlock_5', 'Dense_3', 'kernel'): (1, 768, 768)
+('DiTBlock_5', 'Dense_4', 'bias'): (1, 768)
+('DiTBlock_5', 'Dense_4', 'kernel'): (1, 768, 768)
+('DiTBlock_5', 'MlpBlock_0', 'Dense_0', 'bias'): (1, 3072)
+('DiTBlock_5', 'MlpBlock_0', 'Dense_0', 'kernel'): (1, 768, 3072)
+('DiTBlock_5', 'MlpBlock_0', 'Dense_1', 'bias'): (1, 768)
+('DiTBlock_5', 'MlpBlock_0', 'Dense_1', 'kernel'): (1, 3072, 768)
+('DiTBlock_6', 'Dense_0', 'bias'): (1, 4608)
+('DiTBlock_6', 'Dense_0', 'kernel'): (1, 768, 4608)
+('DiTBlock_6', 'Dense_1', 'bias'): (1, 768)
+('DiTBlock_6', 'Dense_1', 'kernel'): (1, 768, 768)
+('DiTBlock_6', 'Dense_2', 'bias'): (1, 768)
+('DiTBlock_6', 'Dense_2', 'kernel'): (1, 768, 768)
+('DiTBlock_6', 'Dense_3', 'bias'): (1, 768)
+('DiTBlock_6', 'Dense_3', 'kernel'): (1, 768, 768)
+('DiTBlock_6', 'Dense_4', 'bias'): (1, 768)
+('DiTBlock_6', 'Dense_4', 'kernel'): (1, 768, 768)
+('DiTBlock_6', 'MlpBlock_0', 'Dense_0', 'bias'): (1, 3072)
+('DiTBlock_6', 'MlpBlock_0', 'Dense_0', 'kernel'): (1, 768, 3072)
+('DiTBlock_6', 'MlpBlock_0', 'Dense_1', 'bias'): (1, 768)
+('DiTBlock_6', 'MlpBlock_0', 'Dense_1', 'kernel'): (1, 3072, 768)
+('DiTBlock_7', 'Dense_0', 'bias'): (1, 4608)
+('DiTBlock_7', 'Dense_0', 'kernel'): (1, 768, 4608)
+('DiTBlock_7', 'Dense_1', 'bias'): (1, 768)
+('DiTBlock_7', 'Dense_1', 'kernel'): (1, 768, 768)
+('DiTBlock_7', 'Dense_2', 'bias'): (1, 768)
+('DiTBlock_7', 'Dense_2', 'kernel'): (1, 768, 768)
+('DiTBlock_7', 'Dense_3', 'bias'): (1, 768)
+('DiTBlock_7', 'Dense_3', 'kernel'): (1, 768, 768)
+('DiTBlock_7', 'Dense_4', 'bias'): (1, 768)
+('DiTBlock_7', 'Dense_4', 'kernel'): (1, 768, 768)
+('DiTBlock_7', 'MlpBlock_0', 'Dense_0', 'bias'): (1, 3072)
+('DiTBlock_7', 'MlpBlock_0', 'Dense_0', 'kernel'): (1, 768, 3072)
+('DiTBlock_7', 'MlpBlock_0', 'Dense_1', 'bias'): (1, 768)
+('DiTBlock_7', 'MlpBlock_0', 'Dense_1', 'kernel'): (1, 3072, 768)
+('DiTBlock_8', 'Dense_0', 'bias'): (1, 4608)
+('DiTBlock_8', 'Dense_0', 'kernel'): (1, 768, 4608)
+('DiTBlock_8', 'Dense_1', 'bias'): (1, 768)
+('DiTBlock_8', 'Dense_1', 'kernel'): (1, 768, 768)
+('DiTBlock_8', 'Dense_2', 'bias'): (1, 768)
+('DiTBlock_8', 'Dense_2', 'kernel'): (1, 768, 768)
+('DiTBlock_8', 'Dense_3', 'bias'): (1, 768)
+('DiTBlock_8', 'Dense_3', 'kernel'): (1, 768, 768)
+('DiTBlock_8', 'Dense_4', 'bias'): (1, 768)
+('DiTBlock_8', 'Dense_4', 'kernel'): (1, 768, 768)
+('DiTBlock_8', 'MlpBlock_0', 'Dense_0', 'bias'): (1, 3072)
+('DiTBlock_8', 'MlpBlock_0', 'Dense_0', 'kernel'): (1, 768, 3072)
+('DiTBlock_8', 'MlpBlock_0', 'Dense_1', 'bias'): (1, 768)
+('DiTBlock_8', 'MlpBlock_0', 'Dense_1', 'kernel'): (1, 3072, 768)
+('DiTBlock_9', 'Dense_0', 'bias'): (1, 4608)
+('DiTBlock_9', 'Dense_0', 'kernel'): (1, 768, 4608)
+('DiTBlock_9', 'Dense_1', 'bias'): (1, 768)
+('DiTBlock_9', 'Dense_1', 'kernel'): (1, 768, 768)
+('DiTBlock_9', 'Dense_2', 'bias'): (1, 768)
+('DiTBlock_9', 'Dense_2', 'kernel'): (1, 768, 768)
+('DiTBlock_9', 'Dense_3', 'bias'): (1, 768)
+('DiTBlock_9', 'Dense_3', 'kernel'): (1, 768, 768)
+('DiTBlock_9', 'Dense_4', 'bias'): (1, 768)
+('DiTBlock_9', 'Dense_4', 'kernel'): (1, 768, 768)
+('DiTBlock_9', 'MlpBlock_0', 'Dense_0', 'bias'): (1, 3072)
+('DiTBlock_9', 'MlpBlock_0', 'Dense_0', 'kernel'): (1, 768, 3072)
+('DiTBlock_9', 'MlpBlock_0', 'Dense_1', 'bias'): (1, 768)
+('DiTBlock_9', 'MlpBlock_0', 'Dense_1', 'kernel'): (1, 3072, 768)
+('Embed_0', 'embedding'): (1, 256, 1)
+('FinalLayer_0', 'Dense_0', 'bias'): (1, 1536)
+('FinalLayer_0', 'Dense_0', 'kernel'): (1, 768, 1536)
+('FinalLayer_0', 'Dense_1', 'bias'): (1, 16)
+('FinalLayer_0', 'Dense_1', 'kernel'): (1, 768, 16)
+('LabelEmbedder_0', 'Embed_0', 'embedding'): (1, 1001, 768)
+('PatchEmbed_0', 'Conv_0', 'bias'): (1, 768)
+('PatchEmbed_0', 'Conv_0', 'kernel'): (1, 2, 2, 4, 768)
+('TimestepEmbedder_0', 'Dense_0', 'bias'): (1, 768)
+('TimestepEmbedder_0', 'Dense_0', 'kernel'): (1, 256, 768)
+('TimestepEmbedder_0', 'Dense_1', 'bias'): (1, 768)
+('TimestepEmbedder_0', 'Dense_1', 'kernel'): (1, 768, 768)
+('TimestepEmbedder_1', 'Dense_0', 'bias'): (1, 768)
+('TimestepEmbedder_1', 'Dense_0', 'kernel'): (1, 256, 768)
+('TimestepEmbedder_1', 'Dense_1', 'bias'): (1, 768)
+('TimestepEmbedder_1', 'Dense_1', 'kernel'): (1, 768, 768)
+
+ parameter shapes:
+('DiTBlock_0', 'Dense_0', 'bias'): (4608,)
+('DiTBlock_0', 'Dense_0', 'kernel'): (768, 4608)
+('DiTBlock_0', 'Dense_1', 'bias'): (768,)
+('DiTBlock_0', 'Dense_1', 'kernel'): (768, 768)
+('DiTBlock_0', 'Dense_2', 'bias'): (768,)
+('DiTBlock_0', 'Dense_2', 'kernel'): (768, 768)
+('DiTBlock_0', 'Dense_3', 'bias'): (768,)
+('DiTBlock_0', 'Dense_3', 'kernel'): (768, 768)
+('DiTBlock_0', 'Dense_4', 'bias'): (768,)
+('DiTBlock_0', 'Dense_4', 'kernel'): (768, 768)
+('DiTBlock_0', 'MlpBlock_0', 'Dense_0', 'bias'): (3072,)
+('DiTBlock_0', 'MlpBlock_0', 'Dense_0', 'kernel'): (768, 3072)
+('DiTBlock_0', 'MlpBlock_0', 'Dense_1', 'bias'): (768,)
+('DiTBlock_0', 'MlpBlock_0', 'Dense_1', 'kernel'): (3072, 768)
+('DiTBlock_1', 'Dense_0', 'bias'): (4608,)
+('DiTBlock_1', 'Dense_0', 'kernel'): (768, 4608)
+('DiTBlock_1', 'Dense_1', 'bias'): (768,)
+('DiTBlock_1', 'Dense_1', 'kernel'): (768, 768)
+('DiTBlock_1', 'Dense_2', 'bias'): (768,)
+('DiTBlock_1', 'Dense_2', 'kernel'): (768, 768)
+('DiTBlock_1', 'Dense_3', 'bias'): (768,)
+('DiTBlock_1', 'Dense_3', 'kernel'): (768, 768)
+('DiTBlock_1', 'Dense_4', 'bias'): (768,)
+('DiTBlock_1', 'Dense_4', 'kernel'): (768, 768)
+('DiTBlock_1', 'MlpBlock_0', 'Dense_0', 'bias'): (3072,)
+('DiTBlock_1', 'MlpBlock_0', 'Dense_0', 'kernel'): (768, 3072)
+('DiTBlock_1', 'MlpBlock_0', 'Dense_1', 'bias'): (768,)
+('DiTBlock_1', 'MlpBlock_0', 'Dense_1', 'kernel'): (3072, 768)
+('DiTBlock_10', 'Dense_0', 'bias'): (4608,)
+('DiTBlock_10', 'Dense_0', 'kernel'): (768, 4608)
+('DiTBlock_10', 'Dense_1', 'bias'): (768,)
+('DiTBlock_10', 'Dense_1', 'kernel'): (768, 768)
+('DiTBlock_10', 'Dense_2', 'bias'): (768,)
+('DiTBlock_10', 'Dense_2', 'kernel'): (768, 768)
+('DiTBlock_10', 'Dense_3', 'bias'): (768,)
+('DiTBlock_10', 'Dense_3', 'kernel'): (768, 768)
+('DiTBlock_10', 'Dense_4', 'bias'): (768,)
+('DiTBlock_10', 'Dense_4', 'kernel'): (768, 768)
+('DiTBlock_10', 'MlpBlock_0', 'Dense_0', 'bias'): (3072,)
+('DiTBlock_10', 'MlpBlock_0', 'Dense_0', 'kernel'): (768, 3072)
+('DiTBlock_10', 'MlpBlock_0', 'Dense_1', 'bias'): (768,)
+('DiTBlock_10', 'MlpBlock_0', 'Dense_1', 'kernel'): (3072, 768)
+('DiTBlock_11', 'Dense_0', 'bias'): (4608,)
+('DiTBlock_11', 'Dense_0', 'kernel'): (768, 4608)
+('DiTBlock_11', 'Dense_1', 'bias'): (768,)
+('DiTBlock_11', 'Dense_1', 'kernel'): (768, 768)
+('DiTBlock_11', 'Dense_2', 'bias'): (768,)
+('DiTBlock_11', 'Dense_2', 'kernel'): (768, 768)
+('DiTBlock_11', 'Dense_3', 'bias'): (768,)
+('DiTBlock_11', 'Dense_3', 'kernel'): (768, 768)
+('DiTBlock_11', 'Dense_4', 'bias'): (768,)
+('DiTBlock_11', 'Dense_4', 'kernel'): (768, 768)
+('DiTBlock_11', 'MlpBlock_0', 'Dense_0', 'bias'): (3072,)
+('DiTBlock_11', 'MlpBlock_0', 'Dense_0', 'kernel'): (768, 3072)
+('DiTBlock_11', 'MlpBlock_0', 'Dense_1', 'bias'): (768,)
+('DiTBlock_11', 'MlpBlock_0', 'Dense_1', 'kernel'): (3072, 768)
+('DiTBlock_2', 'Dense_0', 'bias'): (4608,)
+('DiTBlock_2', 'Dense_0', 'kernel'): (768, 4608)
+('DiTBlock_2', 'Dense_1', 'bias'): (768,)
+('DiTBlock_2', 'Dense_1', 'kernel'): (768, 768)
+('DiTBlock_2', 'Dense_2', 'bias'): (768,)
+('DiTBlock_2', 'Dense_2', 'kernel'): (768, 768)
+('DiTBlock_2', 'Dense_3', 'bias'): (768,)
+('DiTBlock_2', 'Dense_3', 'kernel'): (768, 768)
+('DiTBlock_2', 'Dense_4', 'bias'): (768,)
+('DiTBlock_2', 'Dense_4', 'kernel'): (768, 768)
+('DiTBlock_2', 'MlpBlock_0', 'Dense_0', 'bias'): (3072,)
+('DiTBlock_2', 'MlpBlock_0', 'Dense_0', 'kernel'): (768, 3072)
+('DiTBlock_2', 'MlpBlock_0', 'Dense_1', 'bias'): (768,)
+('DiTBlock_2', 'MlpBlock_0', 'Dense_1', 'kernel'): (3072, 768)
+('DiTBlock_3', 'Dense_0', 'bias'): (4608,)
+('DiTBlock_3', 'Dense_0', 'kernel'): (768, 4608)
+('DiTBlock_3', 'Dense_1', 'bias'): (768,)
+('DiTBlock_3', 'Dense_1', 'kernel'): (768, 768)
+('DiTBlock_3', 'Dense_2', 'bias'): (768,)
+('DiTBlock_3', 'Dense_2', 'kernel'): (768, 768)
+('DiTBlock_3', 'Dense_3', 'bias'): (768,)
+('DiTBlock_3', 'Dense_3', 'kernel'): (768, 768)
+('DiTBlock_3', 'Dense_4', 'bias'): (768,)
+('DiTBlock_3', 'Dense_4', 'kernel'): (768, 768)
+('DiTBlock_3', 'MlpBlock_0', 'Dense_0', 'bias'): (3072,)
+('DiTBlock_3', 'MlpBlock_0', 'Dense_0', 'kernel'): (768, 3072)
+('DiTBlock_3', 'MlpBlock_0', 'Dense_1', 'bias'): (768,)
+('DiTBlock_3', 'MlpBlock_0', 'Dense_1', 'kernel'): (3072, 768)
+('DiTBlock_4', 'Dense_0', 'bias'): (4608,)
+('DiTBlock_4', 'Dense_0', 'kernel'): (768, 4608)
+('DiTBlock_4', 'Dense_1', 'bias'): (768,)
+('DiTBlock_4', 'Dense_1', 'kernel'): (768, 768)
+('DiTBlock_4', 'Dense_2', 'bias'): (768,)
+('DiTBlock_4', 'Dense_2', 'kernel'): (768, 768)
+('DiTBlock_4', 'Dense_3', 'bias'): (768,)
+('DiTBlock_4', 'Dense_3', 'kernel'): (768, 768)
+('DiTBlock_4', 'Dense_4', 'bias'): (768,)
+('DiTBlock_4', 'Dense_4', 'kernel'): (768, 768)
+('DiTBlock_4', 'MlpBlock_0', 'Dense_0', 'bias'): (3072,)
+('DiTBlock_4', 'MlpBlock_0', 'Dense_0', 'kernel'): (768, 3072)
+('DiTBlock_4', 'MlpBlock_0', 'Dense_1', 'bias'): (768,)
+('DiTBlock_4', 'MlpBlock_0', 'Dense_1', 'kernel'): (3072, 768)
+('DiTBlock_5', 'Dense_0', 'bias'): (4608,)
+('DiTBlock_5', 'Dense_0', 'kernel'): (768, 4608)
+('DiTBlock_5', 'Dense_1', 'bias'): (768,)
+('DiTBlock_5', 'Dense_1', 'kernel'): (768, 768)
+('DiTBlock_5', 'Dense_2', 'bias'): (768,)
+('DiTBlock_5', 'Dense_2', 'kernel'): (768, 768)
+('DiTBlock_5', 'Dense_3', 'bias'): (768,)
+('DiTBlock_5', 'Dense_3', 'kernel'): (768, 768)
+('DiTBlock_5', 'Dense_4', 'bias'): (768,)
+('DiTBlock_5', 'Dense_4', 'kernel'): (768, 768)
+('DiTBlock_5', 'MlpBlock_0', 'Dense_0', 'bias'): (3072,)
+('DiTBlock_5', 'MlpBlock_0', 'Dense_0', 'kernel'): (768, 3072)
+('DiTBlock_5', 'MlpBlock_0', 'Dense_1', 'bias'): (768,)
+('DiTBlock_5', 'MlpBlock_0', 'Dense_1', 'kernel'): (3072, 768)
+('DiTBlock_6', 'Dense_0', 'bias'): (4608,)
+('DiTBlock_6', 'Dense_0', 'kernel'): (768, 4608)
+('DiTBlock_6', 'Dense_1', 'bias'): (768,)
+('DiTBlock_6', 'Dense_1', 'kernel'): (768, 768)
+('DiTBlock_6', 'Dense_2', 'bias'): (768,)
+('DiTBlock_6', 'Dense_2', 'kernel'): (768, 768)
+('DiTBlock_6', 'Dense_3', 'bias'): (768,)
+('DiTBlock_6', 'Dense_3', 'kernel'): (768, 768)
+('DiTBlock_6', 'Dense_4', 'bias'): (768,)
+('DiTBlock_6', 'Dense_4', 'kernel'): (768, 768)
+('DiTBlock_6', 'MlpBlock_0', 'Dense_0', 'bias'): (3072,)
+('DiTBlock_6', 'MlpBlock_0', 'Dense_0', 'kernel'): (768, 3072)
+('DiTBlock_6', 'MlpBlock_0', 'Dense_1', 'bias'): (768,)
+('DiTBlock_6', 'MlpBlock_0', 'Dense_1', 'kernel'): (3072, 768)
+('DiTBlock_7', 'Dense_0', 'bias'): (4608,)
+('DiTBlock_7', 'Dense_0', 'kernel'): (768, 4608)
+('DiTBlock_7', 'Dense_1', 'bias'): (768,)
+('DiTBlock_7', 'Dense_1', 'kernel'): (768, 768)
+('DiTBlock_7', 'Dense_2', 'bias'): (768,)
+('DiTBlock_7', 'Dense_2', 'kernel'): (768, 768)
+('DiTBlock_7', 'Dense_3', 'bias'): (768,)
+('DiTBlock_7', 'Dense_3', 'kernel'): (768, 768)
+('DiTBlock_7', 'Dense_4', 'bias'): (768,)
+('DiTBlock_7', 'Dense_4', 'kernel'): (768, 768)
+('DiTBlock_7', 'MlpBlock_0', 'Dense_0', 'bias'): (3072,)
+('DiTBlock_7', 'MlpBlock_0', 'Dense_0', 'kernel'): (768, 3072)
+('DiTBlock_7', 'MlpBlock_0', 'Dense_1', 'bias'): (768,)
+('DiTBlock_7', 'MlpBlock_0', 'Dense_1', 'kernel'): (3072, 768)
+('DiTBlock_8', 'Dense_0', 'bias'): (4608,)
+('DiTBlock_8', 'Dense_0', 'kernel'): (768, 4608)
+('DiTBlock_8', 'Dense_1', 'bias'): (768,)
+('DiTBlock_8', 'Dense_1', 'kernel'): (768, 768)
+('DiTBlock_8', 'Dense_2', 'bias'): (768,)
+('DiTBlock_8', 'Dense_2', 'kernel'): (768, 768)
+('DiTBlock_8', 'Dense_3', 'bias'): (768,)
+('DiTBlock_8', 'Dense_3', 'kernel'): (768, 768)
+('DiTBlock_8', 'Dense_4', 'bias'): (768,)
+('DiTBlock_8', 'Dense_4', 'kernel'): (768, 768)
+('DiTBlock_8', 'MlpBlock_0', 'Dense_0', 'bias'): (3072,)
+('DiTBlock_8', 'MlpBlock_0', 'Dense_0', 'kernel'): (768, 3072)
+('DiTBlock_8', 'MlpBlock_0', 'Dense_1', 'bias'): (768,)
+('DiTBlock_8', 'MlpBlock_0', 'Dense_1', 'kernel'): (3072, 768)
+('DiTBlock_9', 'Dense_0', 'bias'): (4608,)
+('DiTBlock_9', 'Dense_0', 'kernel'): (768, 4608)
+('DiTBlock_9', 'Dense_1', 'bias'): (768,)
+('DiTBlock_9', 'Dense_1', 'kernel'): (768, 768)
+('DiTBlock_9', 'Dense_2', 'bias'): (768,)
+('DiTBlock_9', 'Dense_2', 'kernel'): (768, 768)
+('DiTBlock_9', 'Dense_3', 'bias'): (768,)
+('DiTBlock_9', 'Dense_3', 'kernel'): (768, 768)
+('DiTBlock_9', 'Dense_4', 'bias'): (768,)
+('DiTBlock_9', 'Dense_4', 'kernel'): (768, 768)
+('DiTBlock_9', 'MlpBlock_0', 'Dense_0', 'bias'): (3072,)
+('DiTBlock_9', 'MlpBlock_0', 'Dense_0', 'kernel'): (768, 3072)
+('DiTBlock_9', 'MlpBlock_0', 'Dense_1', 'bias'): (768,)
+('DiTBlock_9', 'MlpBlock_0', 'Dense_1', 'kernel'): (3072, 768)
+('Embed_0', 'embedding'): (256, 1)
+('FinalLayer_0', 'Dense_0', 'bias'): (1536,)
+('FinalLayer_0', 'Dense_0', 'kernel'): (768, 1536)
+('FinalLayer_0', 'Dense_1', 'bias'): (16,)
+('FinalLayer_0', 'Dense_1', 'kernel'): (768, 16)
+('LabelEmbedder_0', 'Embed_0', 'embedding'): (1001, 768)
+('PatchEmbed_0', 'Conv_0', 'bias'): (768,)
+('PatchEmbed_0', 'Conv_0', 'kernel'): (2, 2, 4, 768)
+('TimestepEmbedder_0', 'Dense_0', 'bias'): (768,)
+('TimestepEmbedder_0', 'Dense_0', 'kernel'): (256, 768)
+('TimestepEmbedder_0', 'Dense_1', 'bias'): (768,)
+('TimestepEmbedder_0', 'Dense_1', 'kernel'): (768, 768)
+('TimestepEmbedder_1', 'Dense_0', 'bias'): (768,)
+('TimestepEmbedder_1', 'Dense_0', 'kernel'): (256, 768)
+('TimestepEmbedder_1', 'Dense_1', 'bias'): (768,)
+('TimestepEmbedder_1', 'Dense_1', 'kernel'): (768, 768)
+┌────────────────────────────────────────────────┐
+│                                                │
+│                                                │
+│                                                │
+│                                                │
+│                  TPU 0,1,2,3                   │
+│                                                │
+│                                                │
+│                                                │
+│                                                │
+└────────────────────────────────────────────────┘
+┌─────────────────────────────────────────────────────────────────────────┐
+│                                                                         │
+│                                                                         │
+│                                                                         │
+│                                                                         │
+│                               TPU 0,1,2,3                               │
+│                                                                         │
+│                                                                         │
+│                                                                         │
+│                                                                         │
+└─────────────────────────────────────────────────────────────────────────┘
+doing the else
+(512, 256, 256, 3)
+encode image shape (128, 256, 256, 3)
+Initializing encoder.
+Incoming encoder shape (128, 256, 256, 3)
+Encoder layer (128, 256, 256, 128)
+doing downsample
+Encoder layer (128, 128, 128, 128)
+doing downsample
+Encoder layer (128, 64, 64, 256)
+doing downsample
+Encoder layer (128, 32, 32, 512)
+Encoder layer (128, 32, 32, 512)
+Encoder layer final (128, 32, 32, 512)
+Encoder layer final (128, 32, 32, 512)
+Final embeddings are size (128, 32, 32, 8)
+After quant (128, 32, 32, 4)
+Calc FID for CFG 1.0 and denoise_timesteps 128
+DiT: Input of shape (512, 32, 32, 4) dtype float32
+DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
+DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
+DiT: Conditioning of shape (512, 768) dtype float32
+z_vectors shape (128, 32, 32, 4)
+Decoder incoming shape (128, 32, 32, 4)
+Decoder input (128, 32, 32, 512)
+Mid Block Decoder layer (128, 32, 32, 512)
+Mid Block Decoder layer (128, 32, 32, 512)
+Decoder layer (128, 64, 64, 512)
+Decoder layer (128, 128, 128, 512)
+Decoder layer (128, 256, 256, 256)
+Decoder layer (128, 256, 256, 128)
+FID is 36.056663513183594
+(512, 256, 256, 3)
+Calc FID for CFG 1.0 and denoise_timesteps 64
+DiT: Input of shape (512, 32, 32, 4) dtype float32
+DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
+DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
+DiT: Conditioning of shape (512, 768) dtype float32
+FID is 36.78165054321289
+(512, 256, 256, 3)
+Calc FID for CFG 1.0 and denoise_timesteps 32
+DiT: Input of shape (512, 32, 32, 4) dtype float32
+DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
+DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
+DiT: Conditioning of shape (512, 768) dtype float32
+FID is 38.99095153808594
+(512, 256, 256, 3)
+Calc FID for CFG 1.0 and denoise_timesteps 16
+DiT: Input of shape (512, 32, 32, 4) dtype float32
+DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
+DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
+DiT: Conditioning of shape (512, 768) dtype float32
+FID is 46.122859954833984
+(512, 256, 256, 3)
+Calc FID for CFG 1.0 and denoise_timesteps 8
+DiT: Input of shape (512, 32, 32, 4) dtype float32
+DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
+DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
+DiT: Conditioning of shape (512, 768) dtype float32
+FID is 69.97744750976562
+(512, 256, 256, 3)
+Calc FID for CFG 1.0 and denoise_timesteps 4
+DiT: Input of shape (512, 32, 32, 4) dtype float32
+DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
+DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
+DiT: Conditioning of shape (512, 768) dtype float32
+FID is 140.8319854736328
+(512, 256, 256, 3)
+Calc FID for CFG 1.0 and denoise_timesteps 2
+DiT: Input of shape (512, 32, 32, 4) dtype float32
+DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
+DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
+DiT: Conditioning of shape (512, 768) dtype float32
+FID is 271.42181396484375
+(512, 256, 256, 3)
+Calc FID for CFG 1.0 and denoise_timesteps 1
+DiT: Input of shape (512, 32, 32, 4) dtype float32
+DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
+DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
+DiT: Conditioning of shape (512, 768) dtype float32
+FID is 264.41375732421875
+(512, 256, 256, 3)
+Calc FID for CFG 1.25 and denoise_timesteps 128
+DiT: Input of shape (512, 32, 32, 4) dtype float32
+DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
+DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
+DiT: Conditioning of shape (512, 768) dtype float32
+FID is 22.577829360961914
+(512, 256, 256, 3)
+Calc FID for CFG 1.25 and denoise_timesteps 64
+DiT: Input of shape (512, 32, 32, 4) dtype float32
+DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
+DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
+DiT: Conditioning of shape (512, 768) dtype float32
+FID is 23.178266525268555
+(512, 256, 256, 3)
+Calc FID for CFG 1.25 and denoise_timesteps 32
+DiT: Input of shape (512, 32, 32, 4) dtype float32
+DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
+DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
+DiT: Conditioning of shape (512, 768) dtype float32
+FID is 25.00735092163086
+(512, 256, 256, 3)
+Calc FID for CFG 1.25 and denoise_timesteps 16
+DiT: Input of shape (512, 32, 32, 4) dtype float32
+DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
+DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
+DiT: Conditioning of shape (512, 768) dtype float32
+FID is 31.010696411132812
+(512, 256, 256, 3)
+Calc FID for CFG 1.25 and denoise_timesteps 8
+DiT: Input of shape (512, 32, 32, 4) dtype float32
+DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
+DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
+DiT: Conditioning of shape (512, 768) dtype float32
+FID is 52.40938949584961
+(512, 256, 256, 3)
+Calc FID for CFG 1.25 and denoise_timesteps 4
+DiT: Input of shape (512, 32, 32, 4) dtype float32
+DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
+DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
+DiT: Conditioning of shape (512, 768) dtype float32
+FID is 119.3035659790039
+(512, 256, 256, 3)
+Calc FID for CFG 1.25 and denoise_timesteps 2
+DiT: Input of shape (512, 32, 32, 4) dtype float32
+DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
+DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
+DiT: Conditioning of shape (512, 768) dtype float32
+FID is 262.1265869140625
+(512, 256, 256, 3)
+Calc FID for CFG 1.25 and denoise_timesteps 1
+DiT: Input of shape (512, 32, 32, 4) dtype float32
+DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
+DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
+DiT: Conditioning of shape (512, 768) dtype float32
+FID is 255.10650634765625
+(512, 256, 256, 3)
+Calc FID for CFG 1.5 and denoise_timesteps 128
+DiT: Input of shape (512, 32, 32, 4) dtype float32
+DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
+DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
+DiT: Conditioning of shape (512, 768) dtype float32
+FID is 14.719727516174316
+(512, 256, 256, 3)
+Calc FID for CFG 1.5 and denoise_timesteps 64
+DiT: Input of shape (512, 32, 32, 4) dtype float32
+DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
+DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
+DiT: Conditioning of shape (512, 768) dtype float32
+FID is 15.140058517456055
+(512, 256, 256, 3)
+Calc FID for CFG 1.5 and denoise_timesteps 32
+DiT: Input of shape (512, 32, 32, 4) dtype float32
+DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
+DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
+DiT: Conditioning of shape (512, 768) dtype float32
+FID is 16.466432571411133
+(512, 256, 256, 3)
+Calc FID for CFG 1.5 and denoise_timesteps 16
+DiT: Input of shape (512, 32, 32, 4) dtype float32
+DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
+DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
+DiT: Conditioning of shape (512, 768) dtype float32
+FID is 21.103317260742188
+(512, 256, 256, 3)
+Calc FID for CFG 1.5 and denoise_timesteps 8
+DiT: Input of shape (512, 32, 32, 4) dtype float32
+DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
+DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
+DiT: Conditioning of shape (512, 768) dtype float32
+FID is 38.82719421386719
+(512, 256, 256, 3)
+Calc FID for CFG 1.5 and denoise_timesteps 4
+DiT: Input of shape (512, 32, 32, 4) dtype float32
+DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
+DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
+DiT: Conditioning of shape (512, 768) dtype float32
+FID is 100.574951171875
+(512, 256, 256, 3)
+Calc FID for CFG 1.5 and denoise_timesteps 2
+DiT: Input of shape (512, 32, 32, 4) dtype float32
+DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
+DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
+DiT: Conditioning of shape (512, 768) dtype float32
+FID is 254.09597778320312
+(512, 256, 256, 3)
+Calc FID for CFG 1.5 and denoise_timesteps 1
+DiT: Input of shape (512, 32, 32, 4) dtype float32
+DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
+DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
+DiT: Conditioning of shape (512, 768) dtype float32
+FID is 248.49542236328125
+(512, 256, 256, 3)
+Calc FID for CFG 1.75 and denoise_timesteps 128
+DiT: Input of shape (512, 32, 32, 4) dtype float32
+DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
+DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
+DiT: Conditioning of shape (512, 768) dtype float32
+FID is 10.582630157470703
+(512, 256, 256, 3)
+Calc FID for CFG 1.75 and denoise_timesteps 64
+DiT: Input of shape (512, 32, 32, 4) dtype float32
+DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
+DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
+DiT: Conditioning of shape (512, 768) dtype float32
+FID is 10.867171287536621
+(512, 256, 256, 3)
+Calc FID for CFG 1.75 and denoise_timesteps 32
+DiT: Input of shape (512, 32, 32, 4) dtype float32
+DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
+DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
+DiT: Conditioning of shape (512, 768) dtype float32
+FID is 11.793257713317871
+(512, 256, 256, 3)
+Calc FID for CFG 1.75 and denoise_timesteps 16
+DiT: Input of shape (512, 32, 32, 4) dtype float32
+DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
+DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
+DiT: Conditioning of shape (512, 768) dtype float32
+FID is 15.195892333984375
+(512, 256, 256, 3)
+Calc FID for CFG 1.75 and denoise_timesteps 8
+DiT: Input of shape (512, 32, 32, 4) dtype float32
+DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
+DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
+DiT: Conditioning of shape (512, 768) dtype float32
+FID is 29.16132926940918
+(512, 256, 256, 3)
+Calc FID for CFG 1.75 and denoise_timesteps 4
+DiT: Input of shape (512, 32, 32, 4) dtype float32
+DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
+DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
+DiT: Conditioning of shape (512, 768) dtype float32
+FID is 84.88493347167969
+(512, 256, 256, 3)
+Calc FID for CFG 1.75 and denoise_timesteps 2
+DiT: Input of shape (512, 32, 32, 4) dtype float32
+DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
+DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
+DiT: Conditioning of shape (512, 768) dtype float32
+FID is 247.1053924560547
+(512, 256, 256, 3)
+Calc FID for CFG 1.75 and denoise_timesteps 1
+DiT: Input of shape (512, 32, 32, 4) dtype float32
+DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
+DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
+DiT: Conditioning of shape (512, 768) dtype float32
+FID is 243.9935302734375
+(512, 256, 256, 3)
+Calc FID for CFG 2.0 and denoise_timesteps 128
+DiT: Input of shape (512, 32, 32, 4) dtype float32
+DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
+DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
+DiT: Conditioning of shape (512, 768) dtype float32
+FID is 8.849721908569336
+(512, 256, 256, 3)
+Calc FID for CFG 2.0 and denoise_timesteps 64
+DiT: Input of shape (512, 32, 32, 4) dtype float32
+DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
+DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
+DiT: Conditioning of shape (512, 768) dtype float32
+FID is 9.026944160461426
+(512, 256, 256, 3)
+Calc FID for CFG 2.0 and denoise_timesteps 32
+DiT: Input of shape (512, 32, 32, 4) dtype float32
+DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
+DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
+DiT: Conditioning of shape (512, 768) dtype float32
+FID is 9.616308212280273
+(512, 256, 256, 3)
+Calc FID for CFG 2.0 and denoise_timesteps 16
+DiT: Input of shape (512, 32, 32, 4) dtype float32
+DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
+DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
+DiT: Conditioning of shape (512, 768) dtype float32
+FID is 12.015625953674316
+(512, 256, 256, 3)
+Calc FID for CFG 2.0 and denoise_timesteps 8
+DiT: Input of shape (512, 32, 32, 4) dtype float32
+DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
+DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
+DiT: Conditioning of shape (512, 768) dtype float32
+FID is 22.730276107788086
+(512, 256, 256, 3)
+Calc FID for CFG 2.0 and denoise_timesteps 4
+DiT: Input of shape (512, 32, 32, 4) dtype float32
+DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
+DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
+DiT: Conditioning of shape (512, 768) dtype float32
+FID is 72.25479125976562
+(512, 256, 256, 3)
+Calc FID for CFG 2.0 and denoise_timesteps 2
+DiT: Input of shape (512, 32, 32, 4) dtype float32
+DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
+DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
+DiT: Conditioning of shape (512, 768) dtype float32
+FID is 241.37213134765625
+(512, 256, 256, 3)
+Calc FID for CFG 2.0 and denoise_timesteps 1
+DiT: Input of shape (512, 32, 32, 4) dtype float32
+DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
+DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
+DiT: Conditioning of shape (512, 768) dtype float32
+FID is 241.18011474609375
+(512, 256, 256, 3)
+Calc FID for CFG 2.25 and denoise_timesteps 128
+DiT: Input of shape (512, 32, 32, 4) dtype float32
+DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
+DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
+DiT: Conditioning of shape (512, 768) dtype float32
+FID is 8.385416030883789
+(512, 256, 256, 3)
+Calc FID for CFG 2.25 and denoise_timesteps 64
+DiT: Input of shape (512, 32, 32, 4) dtype float32
+DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
+DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
+DiT: Conditioning of shape (512, 768) dtype float32
+FID is 8.469977378845215
+(512, 256, 256, 3)
+Calc FID for CFG 2.25 and denoise_timesteps 32
+DiT: Input of shape (512, 32, 32, 4) dtype float32
+DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
+DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
+DiT: Conditioning of shape (512, 768) dtype float32
+FID is 8.870479583740234
+(512, 256, 256, 3)
+Calc FID for CFG 2.25 and denoise_timesteps 16
+DiT: Input of shape (512, 32, 32, 4) dtype float32
+DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
+DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
+DiT: Conditioning of shape (512, 768) dtype float32
+FID is 10.516737937927246
+(512, 256, 256, 3)
+Calc FID for CFG 2.25 and denoise_timesteps 8
+DiT: Input of shape (512, 32, 32, 4) dtype float32
+DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
+DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
+DiT: Conditioning of shape (512, 768) dtype float32
+FID is 18.650209426879883
+(512, 256, 256, 3)
+Calc FID for CFG 2.25 and denoise_timesteps 4
+DiT: Input of shape (512, 32, 32, 4) dtype float32
+DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
+DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
+DiT: Conditioning of shape (512, 768) dtype float32
+FID is 62.06416320800781
+(512, 256, 256, 3)
+Calc FID for CFG 2.25 and denoise_timesteps 2
+DiT: Input of shape (512, 32, 32, 4) dtype float32
+DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
+DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
+DiT: Conditioning of shape (512, 768) dtype float32
+FID is 236.56796264648438
+(512, 256, 256, 3)
+Calc FID for CFG 2.25 and denoise_timesteps 1
+DiT: Input of shape (512, 32, 32, 4) dtype float32
+DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
+DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
+DiT: Conditioning of shape (512, 768) dtype float32
+FID is 238.99615478515625
+(512, 256, 256, 3)
+Calc FID for CFG 2.5 and denoise_timesteps 128
+DiT: Input of shape (512, 32, 32, 4) dtype float32
+DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
+DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
+DiT: Conditioning of shape (512, 768) dtype float32
+FID is 8.667655944824219
+(512, 256, 256, 3)
+Calc FID for CFG 2.5 and denoise_timesteps 64
+DiT: Input of shape (512, 32, 32, 4) dtype float32
+DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
+DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
+DiT: Conditioning of shape (512, 768) dtype float32
+FID is 8.719772338867188
+(512, 256, 256, 3)
+Calc FID for CFG 2.5 and denoise_timesteps 32
+DiT: Input of shape (512, 32, 32, 4) dtype float32
+DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
+DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
+DiT: Conditioning of shape (512, 768) dtype float32
+FID is 8.941210746765137
+(512, 256, 256, 3)
+Calc FID for CFG 2.5 and denoise_timesteps 16
+DiT: Input of shape (512, 32, 32, 4) dtype float32
+DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
+DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
+DiT: Conditioning of shape (512, 768) dtype float32
+FID is 10.057480812072754
+(512, 256, 256, 3)
+Calc FID for CFG 2.5 and denoise_timesteps 8
+DiT: Input of shape (512, 32, 32, 4) dtype float32
+DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
+DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
+DiT: Conditioning of shape (512, 768) dtype float32
+FID is 16.115034103393555
+(512, 256, 256, 3)
+Calc FID for CFG 2.5 and denoise_timesteps 4
+DiT: Input of shape (512, 32, 32, 4) dtype float32
+DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
+DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
+DiT: Conditioning of shape (512, 768) dtype float32
+FID is 53.955543518066406
+(512, 256, 256, 3)
+Calc FID for CFG 2.5 and denoise_timesteps 2
+DiT: Input of shape (512, 32, 32, 4) dtype float32
+DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
+DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
+DiT: Conditioning of shape (512, 768) dtype float32
+FID is 232.41696166992188
+(512, 256, 256, 3)
+Calc FID for CFG 2.5 and denoise_timesteps 1
+DiT: Input of shape (512, 32, 32, 4) dtype float32
+DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
+DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
+DiT: Conditioning of shape (512, 768) dtype float32
+FID is 237.0272216796875
+(512, 256, 256, 3)
+Calc FID for CFG 2.75 and denoise_timesteps 128
+DiT: Input of shape (512, 32, 32, 4) dtype float32
+DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
+DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
+DiT: Conditioning of shape (512, 768) dtype float32
+FID is 9.384163856506348
+(512, 256, 256, 3)
+Calc FID for CFG 2.75 and denoise_timesteps 64
+DiT: Input of shape (512, 32, 32, 4) dtype float32
+DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
+DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
+DiT: Conditioning of shape (512, 768) dtype float32
+FID is 9.374807357788086
+(512, 256, 256, 3)
+Calc FID for CFG 2.75 and denoise_timesteps 32
+DiT: Input of shape (512, 32, 32, 4) dtype float32
+DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
+DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
+DiT: Conditioning of shape (512, 768) dtype float32
+FID is 9.470893859863281
+(512, 256, 256, 3)
+Calc FID for CFG 2.75 and denoise_timesteps 16
+DiT: Input of shape (512, 32, 32, 4) dtype float32
+DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
+DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
+DiT: Conditioning of shape (512, 768) dtype float32
+FID is 10.200791358947754
+(512, 256, 256, 3)
+Calc FID for CFG 2.75 and denoise_timesteps 8
+DiT: Input of shape (512, 32, 32, 4) dtype float32
+DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
+DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
+DiT: Conditioning of shape (512, 768) dtype float32
+FID is 14.65302848815918
+(512, 256, 256, 3)
+Calc FID for CFG 2.75 and denoise_timesteps 4
+DiT: Input of shape (512, 32, 32, 4) dtype float32
+DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
+DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
+DiT: Conditioning of shape (512, 768) dtype float32
+FID is 47.41338348388672
+(512, 256, 256, 3)
+Calc FID for CFG 2.75 and denoise_timesteps 2
+DiT: Input of shape (512, 32, 32, 4) dtype float32
+DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
+DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
+DiT: Conditioning of shape (512, 768) dtype float32
+FID is 228.976318359375
+(512, 256, 256, 3)
+Calc FID for CFG 2.75 and denoise_timesteps 1
+DiT: Input of shape (512, 32, 32, 4) dtype float32
+DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
+DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
+DiT: Conditioning of shape (512, 768) dtype float32
+FID is 235.353271484375
+(512, 256, 256, 3)
+Calc FID for CFG 3.0 and denoise_timesteps 128
+DiT: Input of shape (512, 32, 32, 4) dtype float32
+DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
+DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
+DiT: Conditioning of shape (512, 768) dtype float32
+FID is 10.227725982666016
+(512, 256, 256, 3)
+Calc FID for CFG 3.0 and denoise_timesteps 64
+DiT: Input of shape (512, 32, 32, 4) dtype float32
+DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
+DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
+DiT: Conditioning of shape (512, 768) dtype float32
+FID is 10.197022438049316
+(512, 256, 256, 3)
+Calc FID for CFG 3.0 and denoise_timesteps 32
+DiT: Input of shape (512, 32, 32, 4) dtype float32
+DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
+DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
+DiT: Conditioning of shape (512, 768) dtype float32
+FID is 10.22225570678711
+(512, 256, 256, 3)
+Calc FID for CFG 3.0 and denoise_timesteps 16
+DiT: Input of shape (512, 32, 32, 4) dtype float32
+DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
+DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
+DiT: Conditioning of shape (512, 768) dtype float32
+FID is 10.676050186157227
+(512, 256, 256, 3)
+Calc FID for CFG 3.0 and denoise_timesteps 8
+DiT: Input of shape (512, 32, 32, 4) dtype float32
+DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
+DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
+DiT: Conditioning of shape (512, 768) dtype float32
+FID is 13.905902862548828
+(512, 256, 256, 3)
+Calc FID for CFG 3.0 and denoise_timesteps 4
+DiT: Input of shape (512, 32, 32, 4) dtype float32
+DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
+DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
+DiT: Conditioning of shape (512, 768) dtype float32
+FID is 42.176517486572266
+(512, 256, 256, 3)
+Calc FID for CFG 3.0 and denoise_timesteps 2
+DiT: Input of shape (512, 32, 32, 4) dtype float32
+DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
+DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
+DiT: Conditioning of shape (512, 768) dtype float32
+FID is 226.26129150390625
+(512, 256, 256, 3)
+Calc FID for CFG 3.0 and denoise_timesteps 1
+DiT: Input of shape (512, 32, 32, 4) dtype float32
+DiT: After patch embed, shape is (512, 256, 768) dtype bfloat16
+DiT: Patch Embed of shape (512, 256, 768) dtype bfloat16
+DiT: Conditioning of shape (512, 768) dtype float32
+FID is 233.85470581054688