[feat] attempt at push to hugging face method

uhlmanngroup · Sep 30, 2024 · d30ec91 · d30ec91
1 parent 5458c62
commit d30ec91
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/bioimage_embed/lightning/torch.py b/bioimage_embed/lightning/torch.py
@@ -52,7 +52,10 @@ def __init__(self, model, args=SimpleNamespace()):
         # self.args = SimpleNamespace(**{**merged_kwargs, **vars(self.args)})
         self.save_hyperparameters(vars(self.args))
         # self.model.train()
-
+
+    def push_to_hf_hub(self, *args, **kwargs):
+        return self.model.push_to_hf_hub(*args, **kwargs)
+
     def forward(self, batch):
         x = self.batch_to_tensor(batch)
         return ModelOutput(x=x, out=self.model(x))