diff --git a/src/Tracing.jl b/src/Tracing.jl
index e110407c05..9d15abe6f6 100644
--- a/src/Tracing.jl
+++ b/src/Tracing.jl
@@ -1303,7 +1303,7 @@ Base.@nospecializeinfer function make_tracer(
         end
         res = if toscalar
             TracedRNumber{T}((path,), nothing)
-        elseif tobatch !== nothing
+        elseif tobatch !== nothing && prev.shape != tobatch
             error("This should not happen...")
         else
             TracedRArray{T,N}((path,), prev.mlir_data, size(prev))