diff --git a/fabric/fabric.py b/fabric/fabric.py
index 78f3b11..b1f912a 100644
--- a/fabric/fabric.py
+++ b/fabric/fabric.py
@@ -280,6 +280,7 @@ class FABRICPatcher:
 def get_weights(pos_weight, neg_weight, q, num_pos, num_neg, cond_uncond_idxs):
     """
     Prepare weights for the weighted attention
+    :return: Weights of shape [batch_size, nk] where nk is the size of the Key sequence length. batch_size = len(cond_uncond_idxs)
     """
     input_dim = q.shape[1]
     hs_dim = max(num_pos, num_neg) * input_dim
diff --git a/fabric/weighted_attn.py b/fabric/weighted_attn.py
index 267a9ce..1405388 100644
--- a/fabric/weighted_attn.py
+++ b/fabric/weighted_attn.py
@@ -26,6 +26,7 @@ class Weighted_Attn_Patcher:
                 except Exception as e:
                     print(e)
                     print("[FABRIC] Encountered an exception. If this is not a memory issue, please report this issue.")
+                    print(f"[FABRIC] weights b: {weights.shape[0]}, weights nk: {weights.shape[1]}, nq: {nq}, nk: {nk}, h: {h}, B: {B}")
                     self.unpatch()
                     raise e
 
@@ -45,6 +46,7 @@ class Weighted_Attn_Patcher:
                 except Exception as e:
                     print(e)
                     print("[FABRIC] Encountered an exception. If this is not a memory issue, please report this issue.")
+                    print(f"[FABRIC] weights b: {weights.shape[0]}, weights nk: {weights.shape[1]}, nq: {nq}, nk: {nk}, h: {h}, bs: {bs}")
                     self.unpatch()
                     raise e
             torch.nn.functional.scaled_dot_product_attention = pt_sdp
@@ -81,6 +83,7 @@ class Weighted_Attn_Patcher:
                 except Exception as e:
                     print(e)
                     print("[FABRIC] Encountered an exception. If this is not a memory issue, please report this issue.")
+                    print(f"[FABRIC] weights b: {weights.shape[0]}, weights nk: {weights.shape[1]}, nq: {nq}, nk: {nk}, h: {h}, B: {B}")
                     self.unpatch()
                     raise e
             torch.Tensor.softmax = softmax_method