From 989cdb812dc3392a3cf37e7a5153998a7601498c Mon Sep 17 00:00:00 2001 From: NielsRogge <48327001+NielsRogge@users.noreply.github.com> Date: Sat, 1 Jun 2024 13:05:46 +0200 Subject: [PATCH] Add mixin (#168) --- ultralytics/engine/model.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/ultralytics/engine/model.py b/ultralytics/engine/model.py index ef5c93c0..4339dbd5 100644 --- a/ultralytics/engine/model.py +++ b/ultralytics/engine/model.py @@ -13,8 +13,10 @@ from ultralytics.hub.utils import HUB_WEB_ROOT from ultralytics.nn.tasks import attempt_load_one_weight, guess_model_task, nn, yaml_model_load from ultralytics.utils import ASSETS, DEFAULT_CFG_DICT, LOGGER, RANK, SETTINGS, callbacks, checks, emojis, yaml_load +from huggingface_hub import PyTorchModelHubMixin -class Model(nn.Module): + +class Model(nn.Module, PyTorchModelHubMixin): """ A base class for implementing YOLO models, unifying APIs across different model types.