Update src/TensorStoreWrapper.jl

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

Author: mkitti

src/TensorStoreWrapper.jl

@@ -236,8 +236,9 @@ end
 Base.size(w::IndexDomainWrapper) = pyconvert(Tuple, parent(w).shape)
 function Base.size(w::IndexDomainWrapper, d::Integer)
     d < 1 && throw(ArgumentError("dimension must be ≥ 1"))
+    d > ndims(w) && return 1
     di = Base.to_index(d)
-    return d <= ndims(w) ? pyconvert(Int, parent(w).shape[di-1]) : 1
+    return pyconvert(Int, parent(w).shape[di-1])
 end
 Base.ndims(w::IndexDomainWrapper) = pyconvert(Int, parent(w).rank)
 function Base.axes(w::IndexDomainWrapper)