LAin

pip install -qU "huggingface_hub>=0.22"

from huggingface_hub import PyTorchModelHubMixin
from torch import nn

class MyModel(nn.Module,PyTorchModelHubMixin):
  def __init__(self, a, b):
    super().__init__()
    self.layer = nn.Linear(a,b)
  def forward(self,inputs):
    return self.layer(inputs)

first_model = MyModel(3,1)

first_model.push_to_hub("not-lain/test")

pretrained_model = MyModel.from_pretrained("not-lain/test")