@@ -84,8 +84,8 @@ class MidFlowBlock(nn.Module):
channels, channels, (3, 3, 3), droprate=drop_rate, padding="same"
)
- # self.block = nn.Sequential(self.cnv1, self.cnv2, self.cnv3)
- self.block = self.cnv1
+ self.block = nn.Sequential(self.cnv1, self.cnv2, self.cnv3)
+ #self.block = self.cnv1
def forward(self, x: Float[torch.Tensor, "N C D H W"]):
a = nn.ELU()(self.block(x) + x)