Method __init__

sat/vae_modules/cp_enc_dec.py:532–542

(
        self,
        in_channels,
        with_conv,
        compress_time=False,
    )

class Upsample3D(nn.Module):
    def __init__(
        self,
        in_channels,
        with_conv,
        compress_time=False,
    ):
        super().__init__()
        self.with_conv = with_conv
        if self.with_conv:
            self.conv = torch.nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=1, padding=1)
        self.compress_time = compress_time

    def forward(self, x):
        if self.compress_time and x.shape[2] > 1: