Add print statements

visual.py

@@ -31,15 +31,18 @@ class Attention(nn.Module):
         self.dense = nn.Linear(config.hidden_size, config.hidden_size)
         self.output_dropout = torch.nn.Dropout(config.dropout_prob)
 
-    def forward(self, x: "tensor(B, L, D)") -> "tensor(B, L, D)":
+    def forward(self, x: "tensor(B, L, D)", print_values=False) -> "tensor(B, L, D)":
         B, L, _ = x.shape
         qkv = self.query_key_value(x)
         qkv = qkv.reshape(B, L, 3, self.num_heads, -1).permute(2, 0, 1, 3, 4)  # 3, B, L, H, D
         q, k, v = qkv[0], qkv[1], qkv[2]
 
-        out = xops.memory_efficient_attention(
-            q, k, v, scale=self.scale,
-        )
+        # out = xops.memory_efficient_attention(
+        #     q, k, v, scale=self.scale,
+        # )
+
+        out = self.attention(q, k, v)
+
         output = self.dense(out.view(B, L, -1))
         output = self.output_dropout(output)
         return output
@@ -80,7 +83,7 @@ class TransformerLayer(nn.Module):
         if print_values:
             print("Hidden states before attention:", attention_input[0, :3, :3])
 
-        attention_output = self.attention(attention_input)
+        attention_output = self.attention(attention_input, print_values=print_values)
 
         if print_values:
             print("Hidden states after attention:", attention_output[0, :3, :3])