From 417d1d2ef7fb806dd21500e86e2e496041529ce1 Mon Sep 17 00:00:00 2001 From: Thien Tran Date: Tue, 8 Aug 2023 21:48:09 +0800 Subject: [PATCH] add spacing --- vision_toolbox/backbones/mlp_mixer.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/vision_toolbox/backbones/mlp_mixer.py b/vision_toolbox/backbones/mlp_mixer.py index 600c888..2b1be67 100644 --- a/vision_toolbox/backbones/mlp_mixer.py +++ b/vision_toolbox/backbones/mlp_mixer.py @@ -75,6 +75,7 @@ def from_config(variant: str, patch_size: int, img_size: int, pretrained: bool = # Table 1 in https://arxiv.org/pdf/2105.01601.pdf n_layers, d_model = dict(S=(8, 512), B=(12, 768), L=(24, 1024), H=(32, 1280))[variant] m = MLPMixer(n_layers, d_model, patch_size, img_size) + if pretrained: ckpt = { ("S", 8): "gsam/Mixer-S_8.npz", @@ -86,6 +87,7 @@ def from_config(variant: str, patch_size: int, img_size: int, pretrained: bool = }[(variant, patch_size)] base_url = "https://storage.googleapis.com/mixer_models/" m.load_jax_weights(torch_hub_download(base_url + ckpt)) + return m @torch.no_grad()