Update modeling_quasrav4.py

modeling_quasrav4.py

@@ -136,6 +136,7 @@ class InfinityFormerLayer(nn.Module):
             self.mem_attn_layer_norm = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps)
             self.mem_attn = LinearAttention(config, layer_idx)
         self.ffn = GatedFeedForward(config)
+        self.final_layer_norm = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps)
 
     def forward(self, hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, **kwargs) -> Tuple[torch.Tensor, ...]:
         residual = hidden_states
@@ -150,6 +151,7 @@ class InfinityFormerLayer(nn.Module):
             hidden_states = mem_residual + hidden_states
 
         hidden_states = self.ffn(hidden_states)
+        hidden_states = self.final_layer_norm(hidden_states)
         return (hidden_states,)
 
 class InfinityFormerEmbeddings(nn.Module):
@@ -196,7 +198,6 @@ class InfinityFormerModel(InfinityFormerPreTrainedModel):
         self.config = config
         self.embeddings = InfinityFormerEmbeddings(config)
         self.layers = nn.ModuleList([InfinityFormerLayer(config, layer_idx=i) for i in range(config.num_hidden_layers)])
-        self.final_layer_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
         self.gradient_checkpointing = False
         self.post_init()
 
@@ -215,7 +216,6 @@ class InfinityFormerModel(InfinityFormerPreTrainedModel):
             else:
                 layer_outputs = layer_module(hidden_states, attention_mask=attention_mask)
             hidden_states = layer_outputs[0]
-        hidden_states = self.final_layer_norm(hidden_states)
         if not return_dict:
             return (hidden_states,)
         return BaseModelOutputWithPast(last_hidden_state=hidden_states, past_key_values=None, hidden_states=None, attentions=None)