Fix

Author: co63oc

ppdiffusers/ppdiffusers/models/unet_2d_blocks.py

@@ -720,7 +720,6 @@ def custom_forward(*inputs):
 
                     return custom_forward
 
-                hidden_states = recompute(create_custom_forward(resnet), hidden_states, temb)
                 hidden_states = recompute(
                     create_custom_forward(attn, return_dict=False),
                     hidden_states,
@@ -729,6 +728,7 @@ def custom_forward(*inputs):
                     attention_mask,
                     encoder_attention_mask,
                 )  # [0]
+                hidden_states = recompute(create_custom_forward(resnet), hidden_states, temb, lora_scale)
             else:
                 hidden_states = attn(
                     hidden_states,

ppdiffusers/tests/models/test_models_unet_2d_condition.py

@@ -184,12 +184,12 @@ def test_gradient_checkpointing(self):
         model_2.clear_gradients()
         loss_2 = (out_2 - labels).mean()
         loss_2.backward()
-        # UNetMidBlock2DCrossAttn create_custom_forward increases the difference.
-        self.assertTrue((loss - loss_2).abs() < 1e-03)
+        # UNetMidBlock2DCrossAttn create_custom_forward associates the difference.
+        self.assertTrue((loss - loss_2).abs() < 1e-5)
         named_params = dict(model.named_parameters())
         named_params_2 = dict(model_2.named_parameters())
         for name, param in named_params.items():
-            self.assertTrue(paddle_all_close(param.grad, named_params_2[name].grad, atol=1e-03))
+            self.assertTrue(paddle_all_close(param.grad, named_params_2[name].grad, atol=5e-5))
 
     def test_model_with_attention_head_dim_tuple(self):
         init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common()