Add print statements

visual.py

@@ -83,7 +83,7 @@ class TransformerLayer(nn.Module):
         attention_output = self.attention(attention_input)
 
         if print_values:
-            print("Hidden states before attention:", attention_input[0, :3, :3])
+            print("Hidden states after attention:", attention_output[0, :3, :3])
 
         attention_output = self.input_layernorm(attention_output)
         hidden_states = attention_input + attention_output