Update predictor.py

predictor.py

@@ -76,7 +76,7 @@ def get_points_on_a_grid(
 class CoTrackerOnlinePredictor(torch.nn.Module):
     def __init__(
         self,
-        checkpoint="./checkpoints/scaled_online.pth",
+        checkpoint=None,
         offline=False,
         v2=False,
         window_len=16,
@@ -85,6 +85,15 @@ class CoTrackerOnlinePredictor(torch.nn.Module):
         self.support_grid_size = 6
         model = torch.hub.load("facebookresearch/co-tracker", "cotracker3_online").model
         # build_cotracker(checkpoint, v2=v2, offline=False, window_len=window_len)
+        if checkpoint is not None:
+            with open(checkpoint, "rb") as f:
+                state_dict = torch.load(f, map_location="cpu")
+                if "model" in state_dict:
+                    state_dict = state_dict["model"]
+            model.model.load_state_dict(state_dict)
+            print('LOAD STATE DICT')
+
+        
         self.interp_shape = model.model_resolution
         self.step = model.window_len // 2
         self.model = model