From c31473ed4492fdf26aec4173451f31590021862f Mon Sep 17 00:00:00 2001 From: Pavel Iakubovskii Date: Tue, 28 May 2024 10:41:40 +0000 Subject: [PATCH] Remove float64 cast for OwlVit and OwlV2 to support MPS device (#31071) Remove float64 --- src/transformers/models/owlv2/modeling_owlv2.py | 1 - src/transformers/models/owlvit/modeling_owlvit.py | 1 - 2 files changed, 2 deletions(-) diff --git a/src/transformers/models/owlv2/modeling_owlv2.py b/src/transformers/models/owlv2/modeling_owlv2.py index a7924085fce..05c5cd4595b 100644 --- a/src/transformers/models/owlv2/modeling_owlv2.py +++ b/src/transformers/models/owlv2/modeling_owlv2.py @@ -1276,7 +1276,6 @@ def forward( if query_mask.ndim > 1: query_mask = torch.unsqueeze(query_mask, dim=-2) - pred_logits = pred_logits.to(torch.float64) pred_logits = torch.where(query_mask == 0, -1e6, pred_logits) pred_logits = pred_logits.to(torch.float32) diff --git a/src/transformers/models/owlvit/modeling_owlvit.py b/src/transformers/models/owlvit/modeling_owlvit.py index a7d84455230..ee6d8aa423d 100644 --- a/src/transformers/models/owlvit/modeling_owlvit.py +++ b/src/transformers/models/owlvit/modeling_owlvit.py @@ -1257,7 +1257,6 @@ def forward( if query_mask.ndim > 1: query_mask = torch.unsqueeze(query_mask, dim=-2) - pred_logits = pred_logits.to(torch.float64) pred_logits = torch.where(query_mask == 0, -1e6, pred_logits) pred_logits = pred_logits.to(torch.float32)