bug fix

manipulate_model/model.py

@@ -1,11 +1,12 @@
 import torch
 import torch.nn as nn
+from huggingface_hub import PyTorchModelHubMixin
 
 from manipulate_model.encoder.encoder import Encoder
 from manipulate_model.decoder.decoder import Decoder
 
 
-class Model(nn.Module):
+class Model(nn.Module, PyTorchModelHubMixin):
     def __init__(self, config):
         super(Model, self).__init__()
         self.config = config

manipulate_model/utils.py

@@ -17,7 +17,7 @@ def get_config_and_model(model_root="manipulate_model/demo-model/audio"):
     if isinstance(config.model.decoder, str):
         config.model.decoder = OmegaConf.load(config.model.decoder)
 
-    model = Model(config)
+    model = Model.from_pretrained("arnabdas8901/manipulation_detection_transformer", config)
     #weights = torch.load(os.path.join(model_root, "weights.pt"))
     #model.load_state_dict(weights["model_state_dict"])