Update model.py

model.py

@@ -381,7 +381,11 @@ class StripedHyena(nn.Module):
             x = x * padding_mask[..., None]
 
         for _, block in enumerate(self.blocks):
-            x, _ = block(x, inference_params=None, padding_mask=padding_mask)
+            if self.gradient_checkpointing and self.training:
+                x, _ = self._gradient_checkpointing_func(block.__call__, x, None, padding_mask)
+            else:
+                x, _ = block(x, inference_params=None, padding_mask=padding_mask)
+        
         return x, None
 
     def initialize_inference_params(self):