UICollectionView does not adapt correctly on iPad rotate

I am trying to make an adaptive layout using a UICollectionView

On an iPhone I would like my view to stack, on an iPad I would like it to arrange 3 in a row.

This is working, however when I rotate my iPad and rotate it back again, my subviews do not return the position they were in originally at in that rotation.

You can see the behaviour if you pay attention to the Recognition item in the gif

This is my view controller

class AdaptiveLayoutCollectionView: UICollectionViewController {

    override func viewDidLoad() {
        super.viewDidLoad()

        registerCells()
        configureLayout(with: view.bounds.size)
    }

    override func viewWillTransition(to size: CGSize, with coordinator: UIViewControllerTransitionCoordinator) {
        super.viewWillTransition(to: size, with: coordinator)
        configureLayout(with: size)
    }

    override func traitCollectionDidChange(_ previousTraitCollection: UITraitCollection?) {
        super.traitCollectionDidChange(previousTraitCollection)
        configureLayout(with: view.bounds.size)
    }

    final func setLayout(as layout: CollectionViewLayoutOption) {
        self.layout = layout
        configureLayout(with: view.bounds.size)
    }

    func registerCells() { /* register cells in super view */ }

    private func configureLayout(with size: CGSize) {
        guard let flowLayout = collectionView.collectionViewLayout as? UICollectionViewFlowLayout else { return }

        flowLayout.minimumInteritemSpacing = 0
        flowLayout.minimumLineSpacing = 0
        flowLayout.sectionInset = UIEdgeInsets(top: 8.0, left: 0, bottom: 8.0, right: 0)

        if traitCollection.horizontalSizeClass == .regular {
            //let minItemWidth: CGFloat = 300
            let minItemWidth: CGFloat = size.width / 3
            let numberOfCell = size.width / minItemWidth
            let width = floor((numberOfCell / floor(numberOfCell)) * minItemWidth)
            flowLayout.itemSize = CGSize(width: width, height: 80)
        } else {
            flowLayout.itemSize = CGSize(width: size.width, height: 80)
        }
        collectionView.reloadData()
    }
}