Skip to content

Commit 7fb78a5

Browse files
authored
Add mixin (#469)
1 parent 8ffd905 commit 7fb78a5

File tree

1 file changed

+3
-1
lines changed

1 file changed

+3
-1
lines changed

mamba_ssm/modules/mamba2.py

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -31,8 +31,10 @@
3131
from mamba_ssm.ops.triton.ssd_combined import mamba_chunk_scan_combined
3232
from mamba_ssm.ops.triton.ssd_combined import mamba_split_conv1d_scan_combined
3333

34+
from huggingface_hub import PyTorchModelHubMixin
3435

35-
class Mamba2(nn.Module):
36+
37+
class Mamba2(nn.Module, PyTorchModelHubMixin):
3638
def __init__(
3739
self,
3840
d_model,

0 commit comments

Comments
 (0)