diff --git a/optimum/executorch/attentions/custom_kv_cache.py b/optimum/executorch/attentions/custom_kv_cache.py
index 64b7322..515d1c7 100644
--- a/optimum/executorch/attentions/custom_kv_cache.py
+++ b/optimum/executorch/attentions/custom_kv_cache.py
@@ -174,7 +174,7 @@ def from_legacy_cache(
             dtype = legacy_cache.k_cache.dtype
 
         # assert device is None or device == "cpu"
-        assert dtype is None or dtype == torch.float32
+        assert dtype is None or dtype in (torch.float32, torch.bfloat16, torch.float16)
 
         # Use the legacy cache's max_seq_len if max_cache_len is not specified
         if max_cache_len is None and hasattr(legacy_cache, "max_seq_len"):
diff --git a/optimum/executorch/attentions/custom_sdpa.py b/optimum/executorch/attentions/custom_sdpa.py
index 0f5d0fc..d59e418 100644
--- a/optimum/executorch/attentions/custom_sdpa.py
+++ b/optimum/executorch/attentions/custom_sdpa.py
@@ -81,11 +81,7 @@ def custom_sdpa_with_start_pos_forward(
     key = key.transpose(1, 2)
     value = value.transpose(1, 2)
 
-    # Convert the hell out of the inputs to fp32 and back
     input_dtype = query.dtype
-    query = query.to(torch.float32)
-    key = key.to(torch.float32)
-    value = value.to(torch.float32)
 
     # Ignore the causal flag from kwargs but use the one in module
     kwargs.pop("is_causal", None)