Update app.py

app.py

@@ -163,9 +163,13 @@ def loss_style(prompt, style_embed, style_seed):
 
     # Get token embeddings
     token_embeddings = token_emb_layer(input_ids)
+    
 
     # The new embedding - our special birb word
     replacement_token_embedding = style_embed.to(torch_device)
+    # Assuming token_embeddings has shape [batch_size, seq_length, embedding_dim]
+    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]==6829)] = replacement_token_embedding.to(torch_device)