From b2b6028d3c790d3221bca51a106b3b85f2a876a1 Mon Sep 17 00:00:00 2001 From: HELLORPG Date: Fri, 17 May 2024 11:09:29 +0800 Subject: [PATCH] Fix the issue of mismatch of params for FlashDeformAttnFunction. --- DCNv4_op/DCNv4/modules/flash_deform_attn.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/DCNv4_op/DCNv4/modules/flash_deform_attn.py b/DCNv4_op/DCNv4/modules/flash_deform_attn.py index 0c4922b..d10f72d 100644 --- a/DCNv4_op/DCNv4/modules/flash_deform_attn.py +++ b/DCNv4_op/DCNv4/modules/flash_deform_attn.py @@ -128,13 +128,17 @@ class FlashDeformAttn(nn.Module): raise ValueError( "Last dim of reference_points must be 2 or 4, but get {} instead.".format(reference_points.shape[-1]) ) + + # Cat sampling_offsets and attention_weights, generate sampling_loc_attn: + sampling_locations = sampling_locations.flatten(-3).half() + attention_weights = attention_weights.flatten(-2) + sampling_loc_attn = torch.cat([sampling_locations, attention_weights], dim=-1) output = FlashDeformAttnFunction.apply( value, input_spatial_shapes, input_level_start_index, - sampling_locations, - attention_weights, + sampling_loc_attn, self.im2col_step, self.n_points )