Remove print statements

model/attention.py

@@ -56,10 +56,10 @@ class Attention(nn.Module):
         out = rearrange(
             out, "b nh n hdim -> b n (nh hdim)", nh=self.n_heads, hdim=self.head_dim
         )
-        print(f"\nOutput shape after rearranging: {out.shape} in Attention Block")
+        
 
         out = self.output_proj(out)
-        print(f"Out after projection: {out.shape} in Attention Block\n")
+        
 
         assert out.shape == x.shape, "Output shape should be equal to Input shape"