diff --git a/Sources/MediaEditorHub.swift b/Sources/MediaEditorHub.swift
index 74a6562..8f8f476 100644
--- a/Sources/MediaEditorHub.swift
+++ b/Sources/MediaEditorHub.swift
@@ -198,12 +198,13 @@ class MediaEditorHub: UIViewController {
         thumbsCollectionView.reloadData()
         imagesCollectionView.reloadData()
         thumbsCollectionView.layoutIfNeeded()
+        imagesCollectionView.layoutIfNeeded()
         thumbsCollectionView.selectItem(at: IndexPath(row: selectedThumbIndex, section: 0), animated: false, scrollPosition: .right)
         imagesCollectionView.scrollToItem(at: IndexPath(row: selectedThumbIndex, section: 0), at: .right, animated: false)
     }
 
     private func setupToolbar() {
-        toolbarHeight.constant = isSingleImage ? Constants.doneButtonHeight : Constants.thumbHeight
+        toolbarHeight.constant = isSingleImage ? Constants.toolbarHeight : Constants.thumbHeight
         thumbsCollectionView.isHidden = isSingleImage ? true : false
     }
 
@@ -260,7 +261,7 @@ class MediaEditorHub: UIViewController {
         static var imageCellIdentifier = "imageCell"
         static var capabCellIdentifier = "capabilityCell"
         static var thumbHeight: CGFloat = 64
-        static var doneButtonHeight: CGFloat = 44
+        static var toolbarHeight: CGFloat = 44
     }
 }
 
@@ -324,9 +325,11 @@ extension MediaEditorHub: UICollectionViewDelegateFlowLayout {
     func collectionView(_ collectionView: UICollectionView, layout collectionViewLayout: UICollectionViewLayout, sizeForItemAt indexPath: IndexPath) -> CGSize {
         if collectionView == imagesCollectionView {
             return CGSize(width: imagesCollectionView.frame.width, height: imagesCollectionView.frame.height)
+        } else if collectionView == thumbsCollectionView {
+            return CGSize(width: Constants.thumbHeight, height: Constants.thumbHeight)
+        } else {
+            return CGSize(width: Constants.toolbarHeight, height: Constants.toolbarHeight)
         }
-
-        return CGSize(width: 44, height: 44)
     }
 }