diff --git a/timm/layers/pos_embed.py b/timm/layers/pos_embed.py index ad10578acf..426c1c130b 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