Update app.py

app.py

@@ -259,7 +259,7 @@ def generate_30(
     print('catted shape: ', prompt_embeds.shape)
     pooled_prompt_embeds = torch.cat([pooled_prompt_embeds_a, pooled_prompt_embeds_b])
     print('catted pooled shape: ', pooled_prompt_embeds.shape)
-    pooled_prompt_embeds = torch.mean(pooled_prompt_embeds,dim=0)
+    pooled_prompt_embeds = torch.mean(pooled_prompt_embeds,dim=0,keepdim=True)
     print('meaned pooled shape: ', pooled_prompt_embeds.shape)
 
     # 4. (Optional) Average the pooled embeddings
@@ -271,9 +271,9 @@ def generate_30(
     print('catted shape2: ', prompt_embeds2.shape)
     pooled_prompt_embeds2 = torch.cat([pooled_prompt_embeds_a2, pooled_prompt_embeds_b2])
     print('catted pooled shape 2: ', pooled_prompt_embeds2.shape)
-    pooled_prompt_embeds2 = torch.mean(pooled_prompt_embeds2,dim=1)
+    pooled_prompt_embeds2 = torch.mean(pooled_prompt_embeds2,dim=1,keepdim=True)
     print('pooled meaned shape 2: ', pooled_prompt_embeds2.shape)
-    pooled_prompt_embeds = torch.cat([pooled_prompt_embeds, pooled_prompt_embeds2],dim=2)
+    pooled_prompt_embeds = torch.cat([pooled_prompt_embeds, pooled_prompt_embeds2],dim=1)
     print('catted combined meaned pooled shape: ', pooled_prompt_embeds.shape)
   
     options = {
@@ -388,7 +388,7 @@ def generate_60(
     print('catted shape: ', prompt_embeds.shape)
     pooled_prompt_embeds = torch.cat([pooled_prompt_embeds_a, pooled_prompt_embeds_b])
     print('catted pooled shape: ', pooled_prompt_embeds.shape)
-    pooled_prompt_embeds = torch.mean(pooled_prompt_embeds,dim=0)
+    pooled_prompt_embeds = torch.mean(pooled_prompt_embeds,dim=0,keepdim=True)
     print('meaned pooled shape: ', pooled_prompt_embeds.shape)
 
     # 4. (Optional) Average the pooled embeddings
@@ -400,9 +400,9 @@ def generate_60(
     print('catted shape2: ', prompt_embeds2.shape)
     pooled_prompt_embeds2 = torch.cat([pooled_prompt_embeds_a2, pooled_prompt_embeds_b2])
     print('catted pooled shape 2: ', pooled_prompt_embeds2.shape)
-    pooled_prompt_embeds2 = torch.mean(pooled_prompt_embeds2,dim=1)
+    pooled_prompt_embeds2 = torch.mean(pooled_prompt_embeds2,dim=1,keepdim=True)
     print('pooled meaned shape 2: ', pooled_prompt_embeds2.shape)
-    pooled_prompt_embeds = torch.cat([pooled_prompt_embeds, pooled_prompt_embeds2],dim=2)
+    pooled_prompt_embeds = torch.cat([pooled_prompt_embeds, pooled_prompt_embeds2],dim=1)
     print('catted combined meaned pooled shape: ', pooled_prompt_embeds.shape)
   
     options = {
@@ -517,7 +517,7 @@ def generate_90(
     print('catted shape: ', prompt_embeds.shape)
     pooled_prompt_embeds = torch.cat([pooled_prompt_embeds_a, pooled_prompt_embeds_b])
     print('catted pooled shape: ', pooled_prompt_embeds.shape)
-    pooled_prompt_embeds = torch.mean(pooled_prompt_embeds,dim=0)
+    pooled_prompt_embeds = torch.mean(pooled_prompt_embeds,dim=0,keepdim=True)
     print('meaned pooled shape: ', pooled_prompt_embeds.shape)
 
     # 4. (Optional) Average the pooled embeddings
@@ -529,9 +529,9 @@ def generate_90(
     print('catted shape2: ', prompt_embeds2.shape)
     pooled_prompt_embeds2 = torch.cat([pooled_prompt_embeds_a2, pooled_prompt_embeds_b2])
     print('catted pooled shape 2: ', pooled_prompt_embeds2.shape)
-    pooled_prompt_embeds2 = torch.mean(pooled_prompt_embeds2,dim=1)
+    pooled_prompt_embeds2 = torch.mean(pooled_prompt_embeds2,dim=1,keepdim=True)
     print('pooled meaned shape 2: ', pooled_prompt_embeds2.shape)
-    pooled_prompt_embeds = torch.cat([pooled_prompt_embeds, pooled_prompt_embeds2],dim=2)
+    pooled_prompt_embeds = torch.cat([pooled_prompt_embeds, pooled_prompt_embeds2],dim=1)
     print('catted combined meaned pooled shape: ', pooled_prompt_embeds.shape)
   
     options = {