Update app.py

app.py

@@ -352,7 +352,8 @@ def train_model(epochs):
         
             # Compute image reconstruction loss
             img_loss = criterion(output, original)
-            
+            print(f"Original Prompt Tokens Type: {original_prompt_tokens.dtype}, Shape: {original_prompt_tokens.shape}")
+            print(f"Enhanced Prompt Tokens Type: {enhanced_prompt_tokens.dtype}, Shape: {enhanced_prompt_tokens.shape}")
             # Compute prompt guidance loss (L2 norm between original and enhanced prompt embeddings)
             prompt_loss = torch.norm(original_prompt_tokens - enhanced_prompt_tokens, p=2)