DeformableDETR two stage support bfloat16 (huggingface#30907)

Update modeling_deformable_detr.py
zucchini-nlp · Jun 11, 2024 · e6a72b3 · e6a72b3
1 parent e7af110
commit e6a72b3
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)