Update app.py

app.py

@@ -219,9 +219,8 @@ def generate_30(
     pooled_prompt_embeds_list.append(prompt_embedsb[0])
     prompt_embedsb = prompt_embedsb.hidden_states[-2]
     prompt_embeds_list.append(prompt_embedsb)
-    prompt_embeds = torch.cat(prompt_embeds_list).mean(dim=-1)
-    prompt_embeds = prompt_embeds.repeat(1, 1, 1)
-    pooled_prompt_embeds = torch.cat(pooled_prompt_embeds_list).mean(dim=-1)
+    prompt_embeds = torch.cat(prompt_embeds_list).mean(dim=1, keepdim=True)
+    pooled_prompt_embeds = torch.cat(pooled_prompt_embeds_list).mean(dim=1, keepdim=True)
     
     options = {
         #"prompt": prompt,
@@ -295,9 +294,8 @@ def generate_60(
     pooled_prompt_embeds_list.append(prompt_embedsb[0])
     prompt_embedsb = prompt_embedsb.hidden_states[-2]
     prompt_embeds_list.append(prompt_embedsb)
-    prompt_embeds = torch.cat(prompt_embeds_list).mean(dim=-1)
-    prompt_embeds = prompt_embeds.repeat(1, 1, 1)
-    pooled_prompt_embeds = torch.cat(pooled_prompt_embeds_list).mean(dim=-1)
+    prompt_embeds = torch.cat(prompt_embeds_list).mean(dim=1, keepdim=True)
+    pooled_prompt_embeds = torch.cat(pooled_prompt_embeds_list).mean(dim=1, keepdim=True)
     
     options = {
         #"prompt": prompt,
@@ -371,9 +369,8 @@ def generate_90(
     pooled_prompt_embeds_list.append(prompt_embedsb[0])
     prompt_embedsb = prompt_embedsb.hidden_states[-2]
     prompt_embeds_list.append(prompt_embedsb)
-    prompt_embeds = torch.cat(prompt_embeds_list).mean(dim=-1)
-    prompt_embeds = prompt_embeds.repeat(1, 1, 1)
-    pooled_prompt_embeds = torch.cat(pooled_prompt_embeds_list).mean(dim=-1)
+    prompt_embeds = torch.cat(prompt_embeds_list).mean(dim=1, keepdim=True)
+    pooled_prompt_embeds = torch.cat(pooled_prompt_embeds_list).mean(dim=1, keepdim=True)
     
     options = {
         #"prompt": prompt,