     ///   - firstIndex: The index the calculation needs to start from (the main reason is to just calculate the next cells
     ///   - invalidate: Identifies if the attributes need to be invalidated (reset height etc)
     ///   - updatedAttributes: Callback for the attributes that have been updated
-    ///   - customHeightProvider: Callback that can return a custom height for a given attribute. This is used to update heights with the same identifiers
-    private func createOrUpdateCells(from firstIndex: Int, invalidate: Bool, updatedAttributes: OnAttributesUpdatedHandler?, customHeightProvider: ((attributes: BrickLayoutAttributes) -> CGFloat?)? = nil) {
+    private func createOrUpdateCells(from firstIndex: Int, invalidate: Bool, updatedAttributes: OnAttributesUpdatedHandler?) {
 
         guard let dataSource = dataSource else {
             return

        dataSource.prepareForSizeCalculation(for: brickAttributes, containedIn: width, origin: cellOrigin, invalidate: invalidate, in: self, updatedAttributes: updatedAttributes)

        if let brickFrame = oldOriginalFrame where !invalidate {
            if let customHeight = customHeightProvider?(attributes: brickAttributes) {
                height = customHeight
            } else {
                height = brickFrame.height
            }
            width = brickFrame.width
        } else {
            let size = dataSource.size(for: brickAttributes, containedIn: width, in: self)

        }

        brickAttributes.originalFrame.size.height = height
        createOrUpdateCells(from: index, invalidate: false, updatedAttributes: updatedAttributes)
    }

    func invalidate(at index: Int, updatedAttributes: OnAttributesUpdatedHandler?) {
