add profileing and small conv3d repoducer

2025-02-12 16:29:43 +01:00
parent b279e2dfd7
commit 54a585f2b9
4 changed files with 67 additions and 7 deletions
--- a/vae/unet_causal_3d_blocks.py
+++ b/vae/unet_causal_3d_blocks.py
@ -72,6 +72,7 @@ class CausalConv3d(nn.Module):

        self.conv = nn.Conv3d(chan_in, chan_out, kernel_size,
                              stride=stride, dilation=dilation, **kwargs)
+        print(f"Conv3d: {chan_in}, {chan_out}, {kernel_size}, stride={stride}, dilation={dilation}")

    def forward(self, x):
        x = F.pad(x, self.time_causal_padding, mode=self.pad_mode)