Method get_input

lvdm/models/autoencoder.py:120–128

(self, batch, k)

        return dec, posterior

    def get_input(self, batch, k):
        x = batch[k]
        if x.dim() == 5 and self.input_dim == 4:
            b,c,t,h,w = x.shape
            self.b = b
            self.t = t 
            x = rearrange(x, 'b c t h w -> (b t) c h w')

        return x

    def training_step(self, batch, batch_idx, optimizer_idx):
        inputs = self.get_input(batch, self.image_key)