Update min_dalle/min_dalle.py

min_dalle/min_dalle.py

@@ -41,8 +41,8 @@ class MinDalle:
         self.text_token_count = 64
         #Changed
         self.layer_count = 24 if is_mega else 6
-        self.attention_head_count = 32 if is_mega else 16
-        self.embed_count = 2048 if is_mega else 1024
+        self.attention_head_count = 32 if is_mega else 8
+        self.embed_count = 2048 if is_mega else 512
         self.glu_embed_count = 4096 if is_mega else 2730
         self.text_vocab_count = 50272 if is_mega else 50264
         self.image_vocab_count = 16415 if is_mega else 16384