How to override _beforeTokenTransfer correctly?

Question

I have an ERC721Upgradeable contract that has a predefined _beforeTokenTransfer method.

I want to override it by giving it extra features, like this:

function _beforeTokenTransfer(
    address from,
    address to,
    uint256 tokenId,
    uint256 batchSize
) internal virtual override {
    super._beforeTokenTransfer(from, to, tokenId, batchSize = 1);

    // if transfer is restricted on the contract, we still want to allow burning and minting
    if (!hasRole(TRANSFER_ROLE, address(0)) && from != address(0) && to != address(0)) {
        require(hasRole(TRANSFER_ROLE, from) || hasRole(TRANSFER_ROLE, to), "!TRANSFER_ROLE");
    }

    if(!coAllocatedSmartVault) {
        _setupOwner(to);
        _setupRole(DEFAULT_ADMIN_ROLE, to);
        _setupRole(MINTER_ROLE, to);
        _setupRole(TRANSFER_ROLE, to);
    }
}

The function I am calling to trigger this functionality is safeTransferFrom, which then calls the rest of the functions that are listed below (in order):

function safeTransferFrom(
    address from,
    address to,
    uint256 tokenId
) public virtual override {
    safeTransferFrom(from, to, tokenId, "");
}

function safeTransferFrom(
    address from,
    address to,
    uint256 tokenId,
    bytes memory data
) public virtual override {
    require(_isApprovedOrOwner(_msgSender(), tokenId), "ERC721: caller is not token owner or approved");
    _safeTransfer(from, to, tokenId, data);
}

function _safeTransfer(
    address from,
    address to,
    uint256 tokenId,
    bytes memory data
) internal virtual {
    _transfer(from, to, tokenId);
    require(_checkOnERC721Received(from, to, tokenId, data), "ERC721: transfer to non ERC721Receiver implementer");
}

function _transfer(
    address from,
    address to,
    uint256 tokenId
) internal virtual {
    require(ERC721Upgradeable.ownerOf(tokenId) == from, "ERC721: transfer from incorrect owner");
    require(to != address(0), "ERC721: transfer to the zero address");

    _beforeTokenTransfer(from, to, tokenId, 1);

    // Check that tokenId was not transferred by `_beforeTokenTransfer` hook
    require(ERC721Upgradeable.ownerOf(tokenId) == from, "ERC721: transfer from incorrect owner");

    // Clear approvals from the previous owner
    delete _tokenApprovals[tokenId];

    unchecked {
        // `_balances[from]` cannot overflow for the same reason as described in `_burn`:
        // `from`'s balance is the number of token held, which is at least one before the current
        // transfer.
        // `_balances[to]` could overflow in the conditions described in `_mint`. That would require
        // all 2**256 token ids to be minted, which in practice is impossible.
        _balances[from] -= 1;
        _balances[to] += 1;
    }
    _owners[tokenId] = to;

    emit Transfer(from, to, tokenId);

    _afterTokenTransfer(from, to, tokenId, 1);
}

From my understanding, this should trigger the method I have overridden, but this is not happening.

What am I doing wrong?

EDIT 1:

I have checked that my code is working fine when deploying the smart contract directly. The problem I am having is when I deploy the same smart contract using a factory and a proxy. For some motive I don't know, the ownership of the smart contract is not changing when I use a factory.

This is the line I am pointing at: _setupOwner(to).