From e9d2942e133a3a006cc622aa625551e960061883 Mon Sep 17 00:00:00 2001
From: Niels <niels.rogge1@gmail.com>
Date: Sat, 21 Sep 2024 13:35:10 +0200
Subject: [PATCH] Add mixin

---
 look2hear/models/base_model.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/look2hear/models/base_model.py b/look2hear/models/base_model.py
index 89b36ad..730194d 100644
--- a/look2hear/models/base_model.py
+++ b/look2hear/models/base_model.py
@@ -7,6 +7,8 @@
 import torch
 import torch.nn as nn
 
+from huggingface_hub import PyTorchModelHubMixin
+
 
 def _unsqueeze_to_3d(x):
     """Normalize shape of `x` to [batch, n_chan, time]."""
@@ -32,7 +34,7 @@ def pad_to_appropriate_length(x, lcm):
     return x
 
 
-class BaseModel(nn.Module):
+class BaseModel(nn.Module, PyTorchModelHubMixin, repo_url="https://github.com/JusperLee/Apollo", pipeline_tag="audio-to-audio"):
     def __init__(self, sample_rate, in_chan=1):
         super().__init__()
         self._sample_rate = sample_rate