Warning for notimplemented error and patch_size fix

This commit is contained in:
AdamRashid96 committed 2023-08-03 03:14:31 +00:00
1 parent ffb221a062
commit e40786be08
2 files changed
+8 -5

No files matched your search

@@ -461,7 +461,7 @@ class VanillaDataManager(DataManager, Generic[TDataset]):
def _get_pixel_sampler(self, dataset: TDataset, num_rays_per_batch: int) -> PixelSampler:
"""Infer pixel sampler to use."""
if self.config.patch_size > 1:
if self.config.patch_size > 1 and isinstance(self.config.pixel_sampler, PixelSamplerConfig):
return PatchPixelSamplerConfig().setup(
patch_size=self.config.patch_size, num_rays_per_batch=num_rays_per_batch
)
+7 -4
View File
@@ -319,7 +319,9 @@ class PatchPixelSampler(PixelSampler):
) -> Int[Tensor, "batch_size 3"]:
if isinstance(mask, Tensor):
# Note: if there is a mask, should switch to the base PixelSampler class
raise NotImplementedError()
raise NotImplementedError(
"Masked sampling not implemented for PairPixelSampler. Change Config to PixelSamplerConfig instead."
)
else:
sub_bs = batch_size // (self.config.patch_size**2)
indices = torch.rand((sub_bs, 3), device=device) * torch.tensor(
@@ -379,10 +381,11 @@ class PairPixelSampler(PixelSampler): # pylint: disable=too-few-public-methods
mask: Optional[Tensor] = None,
device: Union[torch.device, str] = "cpu",
) -> Int[Tensor, "batch_size 3"]:
if mask:
if isinstance(mask, Tensor):
# Note: if there is a mask, should switch to the base PixelSampler class
raise NotImplementedError()
raise NotImplementedError(
"Masked sampling not implemented for PairPixelSampler. Change Config to PixelSamplerConfig instead."
)
else:
s = (self.rays_to_sample, 1)
ns = torch.randint(0, num_images, s, dtype=torch.long, device=device)