Help using Python in Julia with PythonCall and DLPack

I am trying to utilize TorchIO on some Julia Arrays. I have gone with the combination of PythonCall.jl and DLPack.jl. Currently, the TorchIO augmentations work when the array is below a specific size, but fall apart after the array gets above a size. This doesn’t happen though when running the same code in pure Python with np arrays. I am not sure how to debug this or if this is an error within PythonCall/DLPack. Below are some screenshots of the code and an attached fully reproducible pluto notebook and colab notebook.

Working Transform Julia (size = 30, 60, 4, 1):

Broken Transform Julia (size = 100, 200, 4, 1):

Working Transform Python (size = 1, 4, 60, 30):
image

Working Transform Python (size = 1, 4, 200, 100):
image

Pluto Notebook:

# ╔═╡ f16ac1c9-1cb8-4abc-b9fb-8ca95706b163
begin
	using Pkg
	Pkg.activate(mktempdir())
	Pkg.add("CondaPkg")
	Pkg.add("PythonCall")
	Pkg.add("DLPack")
	Pkg.add("CairoMakie")
	Pkg.add("PlutoUI")

	using CondaPkg; CondaPkg.add("torchio")
	using PythonCall, DLPack
	using CairoMakie, PlutoUI
end

# ╔═╡ 1ab0cf07-318b-40b1-a311-f1027463fa7a
begin
	tio = pyimport("torchio")
	torch = pyimport("torch")
end

# ╔═╡ d48ca3ff-9910-4c9f-aeb6-2da8807ec902
random_motion = tio.RandomMotion(image_interpolation="bspline")

# ╔═╡ adaeb1c0-8163-4eb1-a062-d3d949be501e
random_noise = tio.RandomNoise()

# ╔═╡ 5e78a0a9-5ce2-483a-bfcc-317aadf27971
function create_circular_mask(h, w, center_circle, radius_circle)
    Y, X = collect(1:h), collect(1:w)'
    dist_from_center = sqrt.((X .- center_circle[1]) .^ 2 .+ (Y .- center_circle[2]) .^ 2)
    mask = dist_from_center .<= radius_circle
    return mask
end

# ╔═╡ 253f181d-2b03-4eb9-b7b6-a088ba387b74
begin
	h, w = 100, 200
	circle = Int.(create_circular_mask(h, w, (w/2, h/2), h/10));
	circle = cat(circle, circle, circle, circle, circle, circle; dims=3)
	circle = reshape(circle, (size(circle)..., 1))
end

# ╔═╡ 27eeb585-a89e-43d6-8561-dacfcfa19e2a
md"""
# Regular Circle
"""

# ╔═╡ 9ef94cb7-b048-4d2c-9fc6-95fcd993bdca
heatmap(circle[:, :, 1, 1], colormap=:grays)

# ╔═╡ 8044b650-a7ac-4403-b80a-67ec63341c5f
md"""
# Noisy Circle
"""

# ╔═╡ 08f7a73c-2f4b-437e-927e-5991a4faca8a
begin
	py_circle = DLPack.share(circle, torch.from_dlpack)
	py_noisy_circle = random_noise(py_circle)
	noisy_circle =  DLPack.wrap(py_noisy_circle, torch.to_dlpack)
	@info size(noisy_circle), py_noisy_circle.shape
end

# ╔═╡ 030154c0-3d63-4430-aedf-7d5dedeb9ab1
heatmap(noisy_circle[:, :, 1, 1], colormap=:grays)

# ╔═╡ 325b1890-037a-4119-8d58-89fe9aae66f8
md"""
# Motion + Noisy Circle
"""

# ╔═╡ 426fb8fd-f4b9-4234-ad1a-a4953039235c
begin
	py_motion_noisy_circle = random_motion(py_noisy_circle)
	motion_noisy_circle =  DLPack.wrap(py_motion_noisy_circle, torch.to_dlpack)
	@info size(motion_noisy_circle), py_motion_noisy_circle.shape
end

# ╔═╡ 1b043b05-4005-4a55-a3e4-e15f10f218db
heatmap(motion_noisy_circle[:, :, 1, 1], colormap=:grays)

Colab Notebook: Google Colab

Surprisingly, the TorchIO tio.RandomNoise() transformation always works in both Julia and Python, regardless of size. The problem is with the combination PythonCall/DLPack and the tio.RandomMotion() transformation

I decided to simply perform this operation in pure Python, saving the numpy array and then loading the array into Julia through NPZ.jl.