Datasets: SampleRays prunes custom fields from Batch

[orperel]

Subclassing WispDataset to, i.e, return a MultiviewBatch with additional supervision fields, and applying SampleRays transform on top may cause some of the fields to be omitted.

For example, in the following, view_idx is omitted:

    def __getitem__(self, idx) -> MultiviewBatch:
        """Retrieve a batch of rays and their corresponding values.
        Rays are precomputed from the dataset's cameras, and are cached within the dataset.
        By default, rays are assumed to have corresponding rgb values, sampled from the dataset's images.

        Returns:
            (MultiviewBatch): A batch of rays and their rgb values. The fields can be accessed as a dictionary:
                "rays" - a wisp.core.Rays pack of ray origins and directions, pre-generated from the dataset camera.
                "rgb" - a torch.Tensor of rgb color which corresponds the gt image's pixel each ray intersects.
                "masks" - a torch.BoolTensor specifying if the ray hits a dense area or not.
                "view_idx" - the index of current camera / view.
                 This is estimated from the alpha channel of the gt image, where mask=True if alpha > 0.5.
        """
        out = MultiviewBatch(
            rays=self.data["rays"][idx],
            rgb=self.data["rgb"][idx],
            masks=self.data["masks"][idx]
            view_idx=idx
        )

        if self.transform is not None:
            out = self.transform(out)

        return out