Spaces:
Sleeping
Sleeping
Update landmark_detection.py
Browse files- landmark_detection.py +2 -2
landmark_detection.py
CHANGED
@@ -56,8 +56,8 @@ def load_model(weights_path):
|
|
56 |
local_dir=".",
|
57 |
)
|
58 |
checkpoint = torch.load(weights_path, map_location=device)
|
59 |
-
model.load_state_dict(checkpoint)
|
60 |
-
|
61 |
model = model.eval()
|
62 |
model = model.to(dtype=dtype)
|
63 |
# model = torch.compile(model, mode="reduce-overhead")
|
|
|
56 |
local_dir=".",
|
57 |
)
|
58 |
checkpoint = torch.load(weights_path, map_location=device)
|
59 |
+
# model.load_state_dict(checkpoint)
|
60 |
+
model.load_state_dict(checkpoint["state_dict_backbone"])
|
61 |
model = model.eval()
|
62 |
model = model.to(dtype=dtype)
|
63 |
# model = torch.compile(model, mode="reduce-overhead")
|