def forward(self, x):
        """
        x: shape = [b, c, t, h, w]

        Returns tuple with x, with new shape = [b, seq_len, c_out]
        """
        x = self.encoder(x)
        b, c, t, h, w = x.shape

        x = x.reshape([b, c, t*h*w])  # flatten spatial / temporal dim
        x = x.permute(0, 2, 1)  # permute the c and seq len for s4

        return x

        x: shape = [b, c, t, h, w]

        Returns tuple with x, with new shape = [b, seq_len, c_out]
        
analyze/convnexts4nd/src/tasks/encoders.py