:construction: wip

config.json

@@ -0,0 +1 @@
+{"input_dim": 784, "hidden_dims": [256, 64, 16, 4, 2]}

model.py

@@ -1,8 +1,11 @@
+from dataclasses import dataclass
 from pathlib import Path
-from typing import Tuple
+from typing import Union, List, Tuple
+
 import torch
-from torch import nn
 import torch.nn.functional as F
+from huggingface_hub import ModelHubMixin
+from torch import nn
 
 
 class Dense(nn.Module):
@@ -42,14 +45,21 @@ class Decoder(nn.Module):
         return self.layers(x)
 
 
-class Autoencoder(nn.Module):
+@dataclass
+class AutoencoderConfig:
+    input_dim: int = 784
+    hidden_dims: Union[Tuple[str], List[str]] = (256, 64, 16, 4, 2)
+
 
-    def __init__(self, input_dim: int = 784, hidden_dims: Tuple[int] = (256, 64, 16, 4, 2)):
+class Autoencoder(nn.Module, ModelHubMixin):
+
+    def __init__(self, config: Union[dict, AutoencoderConfig] = AutoencoderConfig(), **kwargs):
         super().__init__()
-        self.encoder = Encoder(input_dim, *hidden_dims)
-        self.decoder = Decoder(input_dim, *reversed(hidden_dims))
-        self.input_dim = input_dim
-        self.hidden_dims = hidden_dims
+        self.config = AutoencoderConfig(**config) if isinstance(config, dict) else config
+        self.config.__dict__.update(**kwargs)
+
+        self.encoder = Encoder(self.config.input_dim, *self.config.hidden_dims)
+        self.decoder = Decoder(self.config.input_dim, *reversed(self.config.hidden_dims))
 
     def forward(self, x):
         x = x.flatten(1)
@@ -58,6 +68,11 @@ class Autoencoder(nn.Module):
         loss = F.mse_loss(recon, x)
         return recon, latent, loss
 
+    def save_pretrained(self, save_directory, **kwargs):
+        assert 'config' not in kwargs, \
+            "save_pretrained handles passing model config for you, please dont pass it"
+        super().save_pretrained(save_directory, config=self.config.__dict__, **kwargs)
+
 
 class MessageModel:
 

pytorch_model.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e550b61b21a7ffdae47ee0f22740aadf94053d00d129af9e7764a0e547b9f17a
+size 1758947