diff --git a/deepforest/main.py b/deepforest/main.py index a070c4f1..e5350076 100644 --- a/deepforest/main.py +++ b/deepforest/main.py @@ -135,7 +135,7 @@ def load_model(self, model_name="weecology/deepforest-tree", revision='main'): revision (str): The model version ('main', 'v1.0.0', etc.). Returns: - self (object):A trained PyTorch model with its config and weights. + None """ # Load the model using from_pretrained self.create_model() @@ -149,8 +149,6 @@ def load_model(self, model_name="weecology/deepforest-tree", revision='main'): self.label_dict = {"Bird": 0} self.numeric_to_label_dict = {v: k for k, v in self.label_dict.items()} - return self - def use_release(self, check_release=True): """Use the latest DeepForest model release from github and load model. Optionally download if release doesn't exist.