diff --git a/DynamicHeightCollectionView.xcodeproj/project.pbxproj b/DynamicHeightCollectionView.xcodeproj/project.pbxproj
index 250352d..a2e3936 100644
--- a/DynamicHeightCollectionView.xcodeproj/project.pbxproj
+++ b/DynamicHeightCollectionView.xcodeproj/project.pbxproj
@@ -116,6 +116,7 @@
 				TargetAttributes = {
 					BD85B938220DC08B00D33BF9 = {
 						CreatedOnToolsVersion = 10.1;
+						LastSwiftMigration = 1110;
 					};
 				};
 			};
@@ -316,7 +317,7 @@
 				PRODUCT_BUNDLE_IDENTIFIER = com.infoedge.DynamicHeightCollectionView;
 				PRODUCT_NAME = "$(TARGET_NAME)";
 				PROVISIONING_PROFILE_SPECIFIER = "99acres Generic";
-				SWIFT_VERSION = 4.2;
+				SWIFT_VERSION = 5.0;
 				TARGETED_DEVICE_FAMILY = "1,2";
 			};
 			name = Debug;
@@ -335,7 +336,7 @@
 				PRODUCT_BUNDLE_IDENTIFIER = com.infoedge.DynamicHeightCollectionView;
 				PRODUCT_NAME = "$(TARGET_NAME)";
 				PROVISIONING_PROFILE_SPECIFIER = "99acres Generic";
-				SWIFT_VERSION = 4.2;
+				SWIFT_VERSION = 5.0;
 				TARGETED_DEVICE_FAMILY = "1,2";
 			};
 			name = Release;
diff --git a/DynamicHeightCollectionView/Example/TableViewCell.swift b/DynamicHeightCollectionView/Example/TableViewCell.swift
index 79b3ca3..c5f9ca5 100644
--- a/DynamicHeightCollectionView/Example/TableViewCell.swift
+++ b/DynamicHeightCollectionView/Example/TableViewCell.swift
@@ -15,11 +15,11 @@ class TableViewCell: UITableViewCell {
     func configure(with arr: [String]) {
         self.arr = arr
         self.collectionView.reloadData()
-        self.collectionView.layoutIfNeeded()
+       // self.collectionView.layoutIfNeeded()
     }
 }
 
-extension TableViewCell: UICollectionViewDataSource, UICollectionViewDelegateFlowLayout {
+extension TableViewCell: UICollectionViewDataSource {
     func collectionView(_ collectionView: UICollectionView, numberOfItemsInSection section: Int) -> Int {
         return self.arr.count
     }
@@ -29,10 +29,14 @@ extension TableViewCell: UICollectionViewDataSource, UICollectionViewDelegateFlo
         cell.textLabel.text = self.arr[indexPath.row]
         return cell
     }
-    
+}
+
+extension TableViewCell: UICollectionViewDelegateFlowLayout {
     func collectionView(_ collectionView: UICollectionView, layout collectionViewLayout: UICollectionViewLayout, sizeForItemAt indexPath: IndexPath) -> CGSize {
         let text = self.arr[indexPath.row]
-        let cellWidth = text.size(withAttributes:[.font: UIFont.systemFont(ofSize:12.0)]).width + 30.0
-        return CGSize(width: cellWidth, height: 30.0)
+        let textSize = text.size(withAttributes:[.font: UIFont.systemFont(ofSize:12.0)])
+        let cellHeight = textSize.height + 20.0
+        let cellWidth = textSize.width + 30.0
+        return CGSize(width: cellWidth, height: cellHeight)
     }
 }
diff --git a/DynamicHeightCollectionView/Example/ViewController.swift b/DynamicHeightCollectionView/Example/ViewController.swift
index 640a1dc..3c996a7 100644
--- a/DynamicHeightCollectionView/Example/ViewController.swift
+++ b/DynamicHeightCollectionView/Example/ViewController.swift
@@ -9,9 +9,27 @@
 import UIKit
 
 class ViewController: UIViewController {
+    
     @IBOutlet weak var tableView: UITableView!
     
-    var arr = ["Basic Operators", "Strings and Characters", "Collection Types", "Control Flow", "Structures and Classes", "Optional Chaining", "Closures", "Automatic Reference Counting", "Advanced Operators", "Access Control", "Memory Safety", "Generics", "Protocols", "Extensions", "Type Casting", "Nested Types", "Error Handling", "Deinitialization"]
+    var arr = ["Basic Operators",
+               "Strings and Characters",
+               "Collection Types",
+               "Control Flow",
+               "Structures and Classes",
+               "Optional Chaining",
+               "Closures",
+               "Automatic Reference Counting",
+               "Advanced Operators",
+               "Access Control",
+               "Memory Safety",
+               "Generics",
+               "Protocols",
+               "Extensions",
+               "Type Casting",
+               "Nested Types",
+               "Error Handling",
+               "Deinitialization"]
     
     override func viewDidLoad() {
         super.viewDidLoad()
@@ -26,6 +44,7 @@ extension ViewController: UITableViewDataSource, UITableViewDelegate {
     func tableView(_ tableView: UITableView, cellForRowAt indexPath: IndexPath) -> UITableViewCell {
         let cell = tableView.dequeueReusableCell(withIdentifier: "TableViewCell", for: indexPath) as! TableViewCell
         cell.configure(with: self.arr)
+        cell.layoutIfNeeded()
         return cell
     }
 }