Update src/nn/attn.py

src/nn/attn.py

@@ -5,10 +5,6 @@ import einops
 from jaxtyping import Float, Bool
 from torch import Tensor
 from typing import Optional
-from torch.nn.attention.flex_attention import flex_attention
-from matplotlib import pyplot as plt
-
-from pdb import set_trace
 
 class KVCache(nn.Module):
     """