Add torch module hub mixin back in (#250)

openclimatefix · Aug 29, 2024 · 28ed7d8 · 28ed7d8
1 parent eaefc7b
commit 28ed7d8
Showing 1 changed file with 2 additions and 2 deletions.
diff --git 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.
     """