From 1f026b984615ffa47ee3dcb19d197c443f43b106 Mon Sep 17 00:00:00 2001
From: Clement Dumas <clement.dumas@ens-paris-saclay.fr>
Date: Fri, 22 Nov 2024 13:27:43 +0100
Subject: [PATCH] added push to hub option

---
 dictionary_learning/dictionary.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/dictionary_learning/dictionary.py b/dictionary_learning/dictionary.py
index bd16916..1e82d31 100644
--- a/dictionary_learning/dictionary.py
+++ b/dictionary_learning/dictionary.py
@@ -3,15 +3,18 @@
 """
 
 from abc import ABC, abstractclassmethod, abstractmethod
+from huggingface_hub import PyTorchModelHubMixin
+
 import torch as th
 import torch.nn as nn
 import torch.nn.init as init
 from torch.nn.functional import relu
 import einops
 from warnings import warn
+import tempfile
 
 
-class Dictionary(ABC, nn.Module):
+class Dictionary(ABC, nn.Module, PyTorchModelHubMixin):
     """
     A dictionary consists of a collection of vectors, an encoder, and a decoder.
     """