diff --git a/chroma_NAG.py b/chroma_NAG.py
index 59fedb9..a0ca65e 100644
--- a/chroma_NAG.py
+++ b/chroma_NAG.py
@@ -44,7 +44,7 @@ def match_seq_len(tensor_to_resize, reference_tensor):
         return F.pad(tensor_to_resize, padding, "constant", 0)
 
 # This new forward function will contain the NAG logic for Chroma's DoubleStreamBlock.
-def chroma_doublestream_forward_nag(self, img, txt, pe, vec=None, attn_mask=None, distill_vec=None):
+def chroma_doublestream_forward_nag(self, img, txt, pe, vec, attn_mask=None):
     """
     A patched forward function for a DoubleStreamBlock that incorporates Normalized Attention Guidance.
     
@@ -59,8 +59,6 @@ def chroma_doublestream_forward_nag(self, img, txt, pe, vec=None, attn_mask=None
            to create a final, guided attention result for the image tokens.
         d. The rest of the block's operations (MLP, etc.) proceed with this guided result.
     """
-    if distill_vec is not None:
-        vec = distill_vec
     # Deconstruct modulation vectors
     (img_mod1, img_mod2), (txt_mod1, txt_mod2) = vec