Niels преди 10 месеца
родител
ревизия
c835652b37
променени са 1 файла, в които са добавени 3 реда и са изтрити 1 реда
  1. 3 1
      ultralytics/engine/model.py

+ 3 - 1
ultralytics/engine/model.py

@@ -8,6 +8,8 @@ import numpy as np
 import torch
 from PIL import Image
 
+from huggingface_hub import PyTorchModelHubMixin
+
 from ultralytics.cfg import TASK2DATA, get_cfg, get_save_dir
 from ultralytics.engine.results import Results
 from ultralytics.hub import HUB_WEB_ROOT, HUBTrainingSession
@@ -26,7 +28,7 @@ from ultralytics.utils import (
 )
 
 
-class Model(nn.Module):
+class Model(nn.Module, PyTorchModelHubMixin, repo_url="https://github.com/ultralytics/ultralytics", pipeline_tag="object-detection", license="agpl-3.0"):
     """
     A base class for implementing YOLO models, unifying APIs across different model types.