Skip to content

Commit e3dd700

Browse files
committed
update make_image_cvcuda to have default batch dim
1 parent 44db71c commit e3dd700

File tree

1 file changed

+3
-2
lines changed

1 file changed

+3
-2
lines changed

test/common_utils.py

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -400,8 +400,9 @@ def make_image_pil(*args, **kwargs):
400400
return to_pil_image(make_image(*args, **kwargs))
401401

402402

403-
def make_image_cvcuda(*args, **kwargs):
404-
return to_cvcuda_tensor(make_image(*args, **kwargs))
403+
def make_image_cvcuda(*args, batch_dims=(1,), **kwargs):
404+
# explicitly default batch_dims to (1,) since to_cvcuda_tensor requires a batch dimension (ndims == 4)
405+
return to_cvcuda_tensor(make_image(*args, batch_dims=batch_dims, **kwargs))
405406

406407

407408
def make_keypoints(canvas_size=DEFAULT_SIZE, *, num_points=4, dtype=None, device="cpu"):

0 commit comments

Comments
 (0)