From 7c846d9970ed33a0d61747583afe15a68a205d03 Mon Sep 17 00:00:00 2001 From: Ross Wightman Date: Sat, 24 Dec 2022 14:36:29 -0800 Subject: [PATCH] Better vmap compat across recent torch versions --- timm/layers/patch_embed.py | 10 +++++++++- timm/version.py | 2 +- 2 files changed, 10 insertions(+), 2 deletions(-) diff --git a/timm/layers/patch_embed.py b/timm/layers/patch_embed.py index b7416260..764519f2 100644 --- a/timm/layers/patch_embed.py +++ b/timm/layers/patch_embed.py @@ -84,6 +84,14 @@ def resample_patch_embed( Resized patch embedding kernel. """ import numpy as np + try: + import functorch + vmap = functorch.vmap + except ImportError: + if hasattr(torch, 'vmap'): + vmap = torch.vmap + else: + assert False, "functorch or a version of torch with vmap is required for FlexiViT resizing." assert len(patch_embed.shape) == 4, "Four dimensions expected" assert len(new_size) == 2, "New shape should only be hw" @@ -115,7 +123,7 @@ def resample_patch_embed( resampled_kernel = resize_mat_pinv @ kernel.reshape(-1) return resampled_kernel.reshape(new_size) - v_resample_kernel = torch.vmap(torch.vmap(resample_kernel, 0, 0), 1, 1) + v_resample_kernel = vmap(vmap(resample_kernel, 0, 0), 1, 1) return v_resample_kernel(patch_embed) diff --git a/timm/version.py b/timm/version.py index c9cc324d..d1a9131b 100644 --- a/timm/version.py +++ b/timm/version.py @@ -1 +1 @@ -__version__ = '0.8.2dev0' +__version__ = '0.8.3dev0'