Update load_model.py

2024-07-30 22:05:31 +03:00
parent 59de563322
commit 7f8dc9ba67
1 changed files with 2 additions and 0 deletions
@@ -185,5 +185,7 @@ def load_model(model_path, model_cfg_path, segmentor, dtype, device):
        model_class = SAM2VideoPredictor
        model = initialize_model(model_class, model_config, segmentor, image_encoder, memory_attention, memory_encoder, sam_mask_decoder_extra_args, dtype, device)
        model.load_state_dict(sd)
+    else:
+        raise ValueError(f"Segmentor {segmentor} not supported")

    return model