From 4b118ae016d6a4fc28be6438a84e5c2a351f4ea9 Mon Sep 17 00:00:00 2001
From: Niels <niels.rogge1@gmail.com>
Date: Sat, 1 Jun 2024 09:35:37 +0200
Subject: [PATCH] Add mixin

---
 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.