def forward(self, hidden_states, features=None):
    (B, n_patch, hidden) = hidden_states.size()
    (l, h, w) = (((self.img_size[0] // (2 ** self.down_factor)) // self.patch_size[0]), ((self.img_size[1] // (2 ** self.down_factor)) // self.patch_size[1]), ((self.img_size[2] // (2 ** self.down_factor)) // self.patch_size[2]))
    x = hidden_states.permute(0, 2, 1)
    x = x.contiguous().view(B, hidden, l, h, w)
    x = self.conv_more(x)
    for (i, decoder_block) in enumerate(self.blocks):
        if (features is not None):
            skip = (features[i] if (i < self.config.n_skip) else None)
        else:
            skip = None
        x = decoder_block(x, skip=skip)
    return x
def forward(self, hidden_states, features=None):
    (B, n_patch, hidden) = hidden_states.size()
    (l, h, w) = (((self.img_size[0] // (2 ** self.down_factor)) // self.patch_size[0]), ((self.img_size[1] // (2 ** self.down_factor)) // self.patch_size[1]), ((self.img_size[2] // (2 ** self.down_factor)) // self.patch_size[2]))
    x = hidden_states.permute(0, 2, 1)
    x = x.contiguous().view(B, hidden, l, h, w)
    x = self.conv_more(x)
    for (i, decoder_block) in enumerate(self.blocks):
        if (features is not None):
            skip = (features[i] if (i < self.config.n_skip) else None)
        else:
            skip = None
        x = decoder_block(x, skip=skip)
    return x