removed prints

ander-db · Sep 22, 2024 · bad213e · bad213e
1 parent 5415ebd
commit bad213e
Show file tree

Hide file tree

Showing 2 changed files with 0 additions and 5 deletions.
diff --git a/src/unet/UNet_Res_GroupNorm_SiLU_D.py b/src/unet/UNet_Res_GroupNorm_SiLU_D.py
@@ -83,7 +83,6 @@ def _build_time_embeddings(self):
         )
 
     def forward(self, x: torch.Tensor, t: Optional[torch.Tensor]) -> torch.Tensor:
-        print(f"[DEBUG] t: {t}")
         return (
             self._forward_with_time_embeddings(x, t)
             if self.time_embedding_dim is not None
@@ -101,11 +100,8 @@ def _forward_without_time_embeddings(self, g):
         g = self.mid(g)
 
         for i, decoder_layer in enumerate(self.decoder):
-            print(f"[DEBUG] i: {i}")
             attn, attn_map = self.attention[i](g=g, x=skip_connections[-(i + 1)])
-            print(f"[DEBUG] attn.shape: {attn.shape}")
             g = decoder_layer(g=g, x=attn)
-            print(f"[DEBUG] g.shape: {g.shape}")
 
         return self.last_layer(g)
 

diff --git a/tests/test_unet.py b/tests/test_unet.py
@@ -49,7 +49,6 @@ def test_unet_forward_pass(unet_model):
     batch_size = 4
     input_tensor = torch.randn(batch_size, 3, 256, 256)
     output = unet_model(input_tensor, None)
-    print(f"{output.shape=}")
     assert output.shape == (batch_size, 3, 256, 256)