Fix typing error

This commit is contained in:
Jonas Kulhanek committed 2023-06-26 16:29:32 +02:00
1 parent 2e6ac4b10e
commit 62226469ec
1 file changed
+2
@@ -519,6 +519,7 @@ class VanillaDataManager(DataManager, Generic[TDataset]):
@torch.compile
def _next_train_compiled(self, image_batch):
assert self.train_pixel_sampler is not None
batch = self.train_pixel_sampler.sample(image_batch)
ray_indices = batch["indices"]
ray_bundle = self.train_ray_generator(ray_indices)
@@ -534,6 +535,7 @@ class VanillaDataManager(DataManager, Generic[TDataset]):
@torch.compile
def _next_eval_compiled(self, image_batch):
assert self.eval_pixel_sampler is not None
batch = self.eval_pixel_sampler.sample(image_batch)
ray_indices = batch["indices"]
ray_bundle = self.eval_ray_generator(ray_indices)