togethercomputer
/

StripedHyena-Hessian-7B

Text Generation

Model card Files Files and versions Community

Zymrael commited on Dec 16, 2023

Commit

b48efba

•

1 Parent(s): 4313473

chore: add checkpoint import

Files changed (1) hide show

model.py +1 -0

model.py CHANGED Viewed

@@ -6,6 +6,7 @@
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
 from .utils import print_rank_0, column_split
 from .cache import InferenceParams, RecurrentInferenceParams

 import torch
 import torch.nn as nn
 import torch.nn.functional as F
+from torch.utils.checkpoint import checkpoint
 from .utils import print_rank_0, column_split
 from .cache import InferenceParams, RecurrentInferenceParams