Update model.py

model.py

@@ -66,6 +66,7 @@ class AttentionBlock(nn.Module):
         self.inner_mha_cls.rotary_emb.register_buffer("inv_freq", self.inner_mha_cls.rotary_emb.inv_freq)
 
         self.mlp = ParallelGatedMLP(config).to(dtype=mlp_dtype)
+        self.filter_output = None
 
     def forward(self, u, inference_params=None, padding_mask=None, *args, **kwargs):
         if (
@@ -288,6 +289,7 @@ class ParallelGatedConvBlock(nn.Module):
 
         self.proj_norm_fn = self.proj_norm
         self.res_mlp_norm_fn = self.res_mlp_norm
+        self.filter_output = None
 
         if self.config.get("compile", False):
             self.proj_norm_fn = torch.compile(self.proj_norm, fullgraph=True, dynamic=False, mode="reduce-overhead")