Update model.py

model.py

@@ -2,6 +2,7 @@ import torch
 import torch.nn as nn
 import torch.nn.functional as F
 import math
+from huggingface_hub import PyTorchModelHubMixin
 
 
 class RoPEPositionalEncoding(nn.Module):
@@ -173,7 +174,7 @@ class ImprovedBitBlock(nn.Module):
         return x
 
 
-class ImprovedBitNet(nn.Module):
+class ImprovedBitNet(nn.Module, PyTorchModelHubMixin):
     def __init__(
         self,
         vocab_size: int = 30522,