Add print statements

visual.py

@@ -37,11 +37,9 @@ class Attention(nn.Module):
         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 = self.attention(q, k, v)
+        out = xops.memory_efficient_attention(
+            q, k, v, scale=self.scale,
+        )
 
         output = self.dense(out.view(B, L, -1))
         output = self.output_dropout(output)