diff --git a/AKPickerView/AKPickerView.swift b/AKPickerView/AKPickerView.swift
index c8370c2..f9ecbdb 100644
--- a/AKPickerView/AKPickerView.swift
+++ b/AKPickerView/AKPickerView.swift
@@ -165,13 +165,13 @@ private class AKCollectionViewLayout: UICollectionViewFlowLayout {
 
 		return nil
 	}
-
-	private func layoutAttributesForElementsInRect(_ rect: CGRect) -> [AnyObject]? {
+    
+    override func layoutAttributesForElements(in rect: CGRect) -> [UICollectionViewLayoutAttributes]? {
 		switch self.delegate.pickerViewStyleForCollectionViewLayout(self) {
 		case .flat:
 			return super.layoutAttributesForElements(in: rect)
 		case .wheel:
-			var attributes = [AnyObject]()
+			var attributes = [UICollectionViewLayoutAttributes]()
 			if self.collectionView!.numberOfSections > 0 {
 				for i in 0 ..< self.collectionView!.numberOfItems(inSection: 0) {
 					let indexPath = IndexPath(item: i, section: 0)