Update app.py

app.py

@@ -159,7 +159,8 @@ def generate_with_prompt_style(prompt, style, seed = 42):
     token_embeddings = token_emb_layer(input_ids)
     # The new embedding - our special birb word
     replacement_token_embedding = embed[list(embed.keys())[0]].to(torch_device)
-
+    replacement_token_embedding = replacement_token_embedding[:768]  # Adjust the size
+    replacement_token_embedding = replacement_token_embedding.unsqueeze(0)  # Make it [1, 768] if necessary
     # Insert this into the token embeddings
     token_embeddings[0, torch.where(input_ids[0]==338)] = replacement_token_embedding.to(torch_device)