Cells keep getting deallocated while scrolling

[wiencheck]

I noticed flickering in cells during scrolling. It wasn't happening while using default flow layout so I started debugging and found out that while using layout from this library, the cells are getting deallocated randomly and sometimes they just don't appear.
I placed a print statement in the deinit of my custom cell class to catch it

[wiencheck]

@Antondomashnev I've attached a video showing the bug

My code in delegate methods looks like this:

extension NewFavouritesLayoutProvider: AWMozaicLayoutDelegate {
    func prepareMozaicLayout() -> AWMozaicLayout {
        return AWMozaicLayout(delegate: self)
    }

    func collectonView(_ collectionView: UICollectionView, mozaik layout: AWMozaicLayout, geometryInfoFor section: Int) -> AWMozaicLayoutSectionGeometryInfo {
        let numberOfColumns = self.numberOfColumns(in: collectionView)
        let cgNumberOfColumns = CGFloat(numberOfColumns)

        let columnWidth = (collectionView.bounds.width - (2 * sectionInset) - (cgNumberOfColumns - 1) * interItemSpacing) / cgNumberOfColumns
        let columns = Array<AWMozaicLayoutColumn>(repeating: AWMozaicLayoutColumn(width: columnWidth), count: numberOfColumns)
        // Square cells
        let rowHeight = columnWidth

        return AWMozaicLayoutSectionGeometryInfo(rowHeight: rowHeight,
                                                 columns: columns,
                                                 minimumInteritemSpacing: interItemSpacing,
                                                 minimumLineSpacing: lineSpacing,
                                                 sectionInset: UIEdgeInsets(top: sectionInset, left: sectionInset, bottom: sectionInset, right: sectionInset),
                                                 headerHeight: 0,
                                                 footerHeight: 0)
    }

    func collectionView(_ collectionView: UICollectionView, mozaik layout: AWMozaicLayout, mozaikSizeForItemAt indexPath: IndexPath) -> AWMozaicLayoutSize {

        let isLarge = viewModel.itemIsLarge(at: indexPath)
        return AWMozaicLayoutSize(numberOfColumns: isLarge ? 2 : 1, numberOfRows: isLarge ? 2 : 1)
    }
}

Apparently this bug is only happening when cells are sized like in this scenario. If I change the number of columns to 2 and rows to 1 all cells work fine.

Screen.Recording.2020-12-22.at.15.08.38.mov