4
0
mirror of https://github.com/nerfstudio-project/nerfstudio.git synced 2026-11-11 10:34:29 +08:00

fix pixel sampling with masks issue when data is list (#2369)

Co-authored-by: Brent Yi <yibrenth@gmail.com>
This commit is contained in:
anc2001 2023-09-14 09:11:15 -07:00 committed by GitHub
parent e9330d515d
commit 805e6e702d
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

View File

@ -215,7 +215,7 @@ class PixelSampler:
num_rays_in_batch = num_rays_per_batch - (num_images - 1) * num_rays_in_batch
indices = self.sample_method(
num_rays_in_batch, 1, image_height, image_width, mask=batch["mask"][i], device=device
num_rays_in_batch, 1, image_height, image_width, mask=batch["mask"][i].unsqueeze(0), device=device
)
indices[:, 0] = i
all_indices.append(indices)