fix typing issue in dinov2 test
Some checks failed
CI / lint_and_typecheck (push) Has been cancelled
Deploy docs to GitHub Pages / Deploy docs (push) Has been cancelled
Spell checker / Spell check (push) Has been cancelled

This commit is contained in:
Pierre Chapuis 2024-09-18 10:41:25 +02:00
parent 336253f26b
commit 73089a4e2d

View file

@ -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"