DeformableDETR two stage support bfloat16 (#30907)

Update modeling_deformable_detr.py
huggingface · May 20, 2024 · 66b0d9e · 66b0d9e
1 parent 5d0bf59
commit 66b0d9e
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/src/transformers/models/deformable_detr/modeling_deformable_detr.py b/src/transformers/models/deformable_detr/modeling_deformable_detr.py
@@ -1616,8 +1616,8 @@ def gen_encoder_output_proposals(self, enc_output, padding_mask, spatial_shapes)
             valid_width = torch.sum(~mask_flatten_[:, 0, :, 0], 1)
 
             grid_y, grid_x = meshgrid(
-                torch.linspace(0, height - 1, height, dtype=torch.float32, device=enc_output.device),
-                torch.linspace(0, width - 1, width, dtype=torch.float32, device=enc_output.device),
+                torch.linspace(0, height - 1, height, dtype=enc_output.dtype, device=enc_output.device),
+                torch.linspace(0, width - 1, width, dtype=enc_output.dtype, device=enc_output.device),
                 indexing="ij",
             )
             grid = torch.cat([grid_x.unsqueeze(-1), grid_y.unsqueeze(-1)], -1)