Remove print statements

model/patch_embed.py

@@ -10,7 +10,7 @@ def get_patch_positional_embedding(pos_emb_dim, grid_size, device):
     grid_size_h, grid_size_w = grid_size
     grid_h = torch.arange(grid_size_h, dtype=torch.float32, device=device)
     grid_w = torch.arange(grid_size_w, dtype=torch.float32, device=device)
-    grid = torch.meshgrid(grid_h, grid_w,indexing="ij")
+    grid = torch.meshgrid(grid_h, grid_w, indexing="ij")
     grid = torch.stack(grid, dim=0)
 
     grid_h_positions = grid[0].reshape(-1)
@@ -30,8 +30,6 @@ def get_patch_positional_embedding(pos_emb_dim, grid_size, device):
     pos_emb = torch.cat([grid_h_emb, grid_w_emb], dim=-1)
 
     # pos_emb -> (Number of patch tokens, pos_emb_dim)
-    print(f"Grid H emb: {grid_h_emb.shape} in Patch Embedding")
-    print(f"Grid W emb: {grid_w_emb.shape} in Patch Embedding")
 
     return pos_emb
 
@@ -68,7 +66,6 @@ class PatchEmbedding(nn.Module):
         nn.init.constant_(self.patch_embed[0].bias, 0)
 
     def forward(self, x):
-    
         out = rearrange(
             x,
             "b c (nh ph) (nw pw) -> b (nh nw) (ph pw c)",
@@ -76,20 +73,16 @@ class PatchEmbedding(nn.Module):
             pw=self.patch_width,
         )
 
-        print(f"Image shape after rearraging: {out.shape} in Patch Embedding Layer")
         actual_h = x.shape[2]  # Height from input tensor
         actual_w = x.shape[3]  # Width from input tensor
         grid_size_h = actual_h // self.patch_height
         grid_size_w = actual_w // self.patch_width
-        
-        
         out = self.patch_embed(out)
         pos_emb = get_patch_positional_embedding(
             self.hidden_dim, grid_size=(grid_size_h, grid_size_w), device=x.device
         )
 
         out += pos_emb
-        print(f"Patch Embeddings: {out.shape}\n")
         return out
 
 

model/transformer.py

@@ -93,8 +93,6 @@ class DIT(nn.Module):
             1 + pre_mlp_scale.unsqueeze(1)
         ) + pre_mlp_shift.unsqueeze(1)
 
-        print(f"\nOutput before unpatchify block {out.shape} in DIT block")
-
         actual_h = x.shape[2]  # Height from input tensor
         actual_w = x.shape[3]  # Width from input tensor
         actual_nh = actual_h // self.patch_height
@@ -102,7 +100,6 @@ class DIT(nn.Module):
 
         # Unpatichify
         out = self.out_proj(out)
-        print(f"Output after projection: {out.shape} in DIT block")
         out = rearrange(
             out,
             "b (nh nw) (ph pw c) -> b c (nh ph) (nw pw)",
@@ -112,7 +109,6 @@ class DIT(nn.Module):
             nh=actual_nh,
         )
 
-        print(f"Output after unpatchify block: {out.shape} in DIT block\n")
         return out