From 73089a4e2d9f98ad7b6c66b09d28471090c13dda Mon Sep 17 00:00:00 2001 From: Pierre Chapuis Date: Wed, 18 Sep 2024 10:41:25 +0200 Subject: [PATCH] fix typing issue in dinov2 test --- tests/foundationals/dinov2/test_dinov2.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/tests/foundationals/dinov2/test_dinov2.py b/tests/foundationals/dinov2/test_dinov2.py index d994d57..47f7959 100644 --- a/tests/foundationals/dinov2/test_dinov2.py +++ b/tests/foundationals/dinov2/test_dinov2.py @@ -60,13 +60,14 @@ def ref_model( if "reg" not in flavor: kwargs["interpolate_offset"] = 0.0 - model = torch.hub.load( # type: ignore + model: torch.nn.Module = torch.hub.load( # type: ignore model=flavor, repo_or_dir=str(dinov2_repo_path), source="local", pretrained=False, # to turn off automatic weights download (see load_state_dict below) **kwargs, - ).to(device=test_device) + ) + model = model.to(device=test_device) flavor = flavor.replace("_reg", "_reg4") weights = test_weights_path / f"{flavor}_pretrain.pth"