             return;
         }
         
-        assert(delegate.conformsToProtocol(CollectionViewWaterfallLayoutDelegate), "UICollectionView's delegate should conform to WaterfallLayoutDelegate protocol")
+        assert(delegate?.conformsToProtocol(CollectionViewWaterfallLayoutDelegate), "UICollectionView's delegate should conform to WaterfallLayoutDelegate protocol")
         assert(columnCount > 0, "WaterfallFlowLayout's columnCount should be greater than 0")
         
         // Initialize variables

    }
    
    //MARK: Private Properties
    private var delegate: CollectionViewWaterfallLayoutDelegate  {
        get {
            return collectionView.delegate as CollectionViewWaterfallLayoutDelegate
        }
    }
    private var columnHeights = [Float]()

            * 2. Section header
            */
            var headerHeight: Float
            if let height = delegate?.collectionView?(collectionView, layout: self, heightForHeaderInSection: section) {
                headerHeight = height
            }
            else {
