Update chatNT.py

chatNT.py

@@ -1330,6 +1330,9 @@ class MultiHeadAttention(nn.Module):
             )
         else:
             attention_weights = F.softmax(attention_weights, dim=-1)
+
+        print(f"Attention weights : {attention_weights.dtype}")
+        print(f"Value heads  : {value_heads.dtype}")
         value_out = torch.einsum(
             "...htT, ...Thd->...thd", attention_weights, value_heads
         )