Update modeling_llama_nsa.py

modeling_llama_nsa.py

@@ -307,7 +307,7 @@ class LlamaNSAAttention(nn.Module):
                 window_size=self.config.window_size,
                 head_first=False,
             )
-        #attn_weights = None
+            attn_weights = None
         
         sa_loss = 0 #torch.nn.SmoothL1Loss()(attn_output_mha, attn_output.detach()) + torch.nn.SmoothL1Loss()(attn_output_mha.detach(), attn_output)