Update app.py

app.py

@@ -226,6 +226,7 @@ def generate_30(
     pooled_prompt_embeds = torch.cat([pooled_prompt_embeds_a, pooled_prompt_embeds_b], dim=1)
 
     # 4. (Optional) Average the pooled embeddings
+    prompt_embeds = prompt_embeds.mean(dim=1, keepdim=True)
     pooled_prompt_embeds = pooled_prompt_embeds.mean(dim=1, keepdim=True)
 
 
@@ -308,6 +309,7 @@ def generate_60(
     pooled_prompt_embeds = torch.cat([pooled_prompt_embeds_a, pooled_prompt_embeds_b], dim=1)
 
     # 4. (Optional) Average the pooled embeddings
+    prompt_embeds = prompt_embeds.mean(dim=1, keepdim=True)
     pooled_prompt_embeds = pooled_prompt_embeds.mean(dim=1, keepdim=True)
 
 
@@ -390,6 +392,7 @@ def generate_90(
     pooled_prompt_embeds = torch.cat([pooled_prompt_embeds_a, pooled_prompt_embeds_b], dim=1)
 
     # 4. (Optional) Average the pooled embeddings
+    prompt_embeds = prompt_embeds.mean(dim=1, keepdim=True)
     pooled_prompt_embeds = pooled_prompt_embeds.mean(dim=1, keepdim=True)