import torch
import torch.nn as nn
from huggingface_hub import PyTorchModelHubMixin
from model import BERT

# create model

model = BERT()
# save locally
model.save_pretrained("genre-pred")

# push to the hub
model.push_to_hub("genre-pred")

# reload
model = MyModel.from_pretrained("rishikasrinivas/genre-pred")