Unexpected error on `report_package`

[albertomercurio]

I have the following definition in my function

using LinearAlgebra

function thermal_dm(N::Int, n::T) where {T<:Real}
    P = rand(Complex{T}, N)
    Z = sum(P)
    P ./= Z
    return __get_size(diagm(0 => P))
end

__get_size(A::Matrix) = size(A)

But when I run

JET.report_package(QuantumToolbox; target_defined_modules = true, ignore_missing_comparison = true)

I get

═════ 1 possible error found ═════
┌ thermal_dm(N::Int64, n::T) where T<:Real @ QuantumToolbox /home/alberto/.julia/dev/QuantumToolbox/src/qobj/states.jl:137
│ no matching method found `__get_size(::Array{Float64, 3})` (1/2 union split): QuantumToolbox.__get_size(QuantumToolbox.diagm(0 QuantumToolbox.:(=>) P::Pair{Int64, Vector{T}} where T<:Complex)::Union{Array{Float64, 3}, Matrix})
└────────────────────

Most importantly, it works with

function thermal_dm(N::Int, n::T) where {T<:Real}
    P = rand(Complex{T}, N)
    Z = sum(P)
    return __get_size(diagm(0 => P ./ Z))
end