diff --git a/tests/e2e/test_diffusion.py b/tests/e2e/test_diffusion.py
index 0c4f71e..23582e8 100644
--- a/tests/e2e/test_diffusion.py
+++ b/tests/e2e/test_diffusion.py
@@ -1395,7 +1395,8 @@ def test_diffusion_sdxl_lora(
     prompt = "professional portrait photo of a girl, photograph, highly detailed face, depth of field, moody light, golden hour, style by Dan Winters, Russell James, Steve McCurry, centered, extremely detailed, Nikon D850, award winning photography"
     negative_prompt = "3d render, cartoon, drawing, art, low light, blur, pixelated, low resolution, black and white"
 
-    SDLoraManager(sdxl).add_loras("dpo", lora_weights, scale=lora_scale)
+    manager = SDLoraManager(sdxl, unet_inclusions=["CrossAttentionBlock"])
+    manager.add_loras("dpo", lora_weights, scale=lora_scale)
 
     clip_text_embedding, pooled_text_embedding = sdxl.compute_clip_text_embedding(
         text=prompt, negative_text=negative_prompt
diff --git a/tests/e2e/test_diffusion_ref/expected_sdxl_dpo_lora.png b/tests/e2e/test_diffusion_ref/expected_sdxl_dpo_lora.png
index 3225508..7a4d824 100644
Binary files a/tests/e2e/test_diffusion_ref/expected_sdxl_dpo_lora.png and b/tests/e2e/test_diffusion_ref/expected_sdxl_dpo_lora.png differ