I have a dataset of RGB tiffs (3 channel) and a dataset of depth tiffs (1 channel) I am dstack-ing them in the getitem method of my dataset class as shown here
    def __getitem__(self, i):
        id = self.ids[i]
        image_path = os.path.join(self.images_dir, id)
        dsm_image_path = os.path.join(self.dsm_images_dir, id)
        mask_path = os.path.join(self.masks_dir, id)
        loss_mask_path = os.path.join(self.loss_masks_dir, id) if self.loss_masks_dir else None # loss masks should have same filename as og image

        # print(id)
        # concat 3 RGB and 1 Depth channel
        rgb_image = self.read_image(image_path)
        print('Converting from',rgb_image.shape)

        dsm_image = self.read_image(dsm_image_path)
        # TODO To optimize training slightly, this can be done as a pre-processing step instead of during training
        # resize dsm image to rgb image if sizes are different
        if dsm_image.shape[:2] != rgb_image.shape[:2]:
            rgb_size = tuple(rgb_image.shape[:2])
            print('resizing', dsm_image.shape, 'to', rgb_image.shape)
            dsm_image = np.array(Image.fromarray(dsm_image).resize(rgb_size)) # resize (bicubic interpol.), change back to (H,W,1) dims
            dsm_image = dsm_image[:,:, None]
            print('dsm', dsm_image.shape)
        # print(rgb_image.shape)
        image_full = np.dstack([rgb_image, self.normalize_image(dsm_image)])
        print(image_full.shape)
        # raise(Exception)
        # read data sample
        sample = dict(
            id=id,
            image=image_full,
            mask=self.read_mask(mask_path),
        )
        if loss_mask_path:
            sample['loss_mask'] = self.read_mask(loss_mask_path)
            sample["loss_mask"] = sample["loss_mask"][None] # expand first dim for loss mask

        # apply augmentations, loss_mask is also augmented. May want to change that
        if self.transform is not None:
            sample = self.transform(**sample)

        # print(image_full.shape)
        sample["mask"] = sample["mask"][None]  # expand first dim for mask, ex size(3, 768, 768) -> size(3, 1, 768, 768)
        return sample