From 40a518c194b406a94c285cc4ca9a03de37235132 Mon Sep 17 00:00:00 2001
From: Sepehr Sameni <Sepehr.Sameni@gmail.com>
Date: Fri, 28 Jul 2023 16:01:42 -0700
Subject: [PATCH] use float in resample_abs_pos_embed_nhwc

since F.interpolate doesn't always support BFloat16
---
 timm/layers/pos_embed.py | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/timm/layers/pos_embed.py b/timm/layers/pos_embed.py
index ad10578a..426c1c13 100644
--- a/timm/layers/pos_embed.py
+++ b/timm/layers/pos_embed.py
@@ -64,12 +64,14 @@ def resample_abs_pos_embed_nhwc(
     if new_size[0] == posemb.shape[-3] and new_size[1] == posemb.shape[-2]:
         return posemb
 
+    previous_dtype = posemb.dtype
+    posemb = posemb.float()
     # do the interpolation
     posemb = posemb.reshape(1, posemb.shape[-3], posemb.shape[-2], posemb.shape[-1]).permute(0, 3, 1, 2)
     posemb = F.interpolate(posemb, size=new_size, mode=interpolation, antialias=antialias)
-    posemb = posemb.permute(0, 2, 3, 1)
+    posemb = posemb.permute(0, 2, 3, 1).to(previous_dtype)
 
     if not torch.jit.is_scripting() and verbose:
         _logger.info(f'Resized position embedding: {posemb.shape[-3:-1]} to {new_size}.')
 
-    return posemb
\ No newline at end of file
+    return posemb