From 28ed7d823207f9ca4c069e4f73f9b82e94a4c04f Mon Sep 17 00:00:00 2001
From: Sukhil Patel <42407101+Sukh-P@users.noreply.github.com>
Date: Thu, 29 Aug 2024 09:40:26 +0100
Subject: [PATCH] Add torch module hub mixin back in (#250)

---
 pvnet/models/base_model.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/pvnet/models/base_model.py b/pvnet/models/base_model.py
index 081e9bc5..9faff47d 100644
--- a/pvnet/models/base_model.py
+++ b/pvnet/models/base_model.py
@@ -13,7 +13,7 @@
 import torch.nn.functional as F
 import wandb
 import yaml
-from huggingface_hub import ModelCard, ModelCardData
+from huggingface_hub import ModelCard, ModelCardData, PyTorchModelHubMixin
 from huggingface_hub.constants import CONFIG_NAME, PYTORCH_WEIGHTS_NAME
 from huggingface_hub.file_download import hf_hub_download
 from huggingface_hub.hf_api import HfApi
@@ -140,7 +140,7 @@ def minimize_data_config(input_path, output_path, model):
         yaml.dump(config, outfile, default_flow_style=False)
 
 
-class PVNetModelHubMixin:
+class PVNetModelHubMixin(PyTorchModelHubMixin):
     """
     Implementation of [`PyTorchModelHubMixin`] to provide model Hub upload/download capabilities.
     """