def train_on_batch(self, batch, train_step):
        """
        Expects a batch of images and depth as input
        batch["image"].shape : batch_size, c, h, w
        batch["depth"].shape : batch_size, 1, h, w

        Assumes all images in a batch are from the same dataset
        """

        images, depths_gt = batch['image'].to(
            self.device), batch['depth'].to(self.device)
        # batch['dataset'] is a tensor strings all valued either 'nyu' or 'kitti'. labels nyu -> 0, kitti -> 1
        dataset = batch['dataset'][0]
        # Convert to 0s or 1s
        domain_labels = torch.Tensor([dataset == 'kitti' for _ in range(
            images.size(0))]).to(torch.long).to(self.device)

        # m = self.model.module if self.config.multigpu else self.model

        b, c, h, w = images.size()
        mask = batch["mask"].to(self.device).to(torch.bool)

        losses = {}

        with amp.autocast(enabled=self.config.use_amp):
            output = self.model(images)
            pred_depths = output['metric_depth']
            domain_logits = output['domain_logits']

            l_si, pred = self.silog_loss(
                pred_depths, depths_gt, mask=mask, interpolate=True, return_interpolated=True)
            loss = self.config.w_si * l_si
            losses[self.silog_loss.name] = l_si

            if self.config.w_grad > 0:
                l_grad = self.grad_loss(pred, depths_gt, mask=mask)
                loss = loss + self.config.w_grad * l_grad
                losses[self.grad_loss.name] = l_grad
            else:
                l_grad = torch.Tensor([0])

            if self.config.w_domain > 0:
                l_domain = self.domain_classifier_loss(
                    domain_logits, domain_labels)
                loss = loss + self.config.w_domain * l_domain
                losses["DomainLoss"] = l_domain
            else:
                l_domain = torch.Tensor([0.])

        self.scaler.scale(loss).backward()

        if self.config.clip_grad > 0:
            self.scaler.unscale_(self.optimizer)
            nn.utils.clip_grad_norm_(
                self.model.parameters(), self.config.clip_grad)

        self.scaler.step(self.optimizer)

        if self.should_log and self.step > 1 and (self.step % int(self.config.log_images_every * self.iters_per_epoch)) == 0:
            depths_gt[torch.logical_not(mask)] = -99
            self.log_images(rgb={"Input": images[0, ...]}, depth={"GT": depths_gt[0], "PredictedMono": pred[0]}, prefix="Train",
                            min_depth=DATASETS_CONFIG[dataset]['min_depth'], max_depth=DATASETS_CONFIG[dataset]['max_depth'])

        self.scaler.update()
        self.optimizer.zero_grad(set_to_none=True)

        return losses

        Expects a batch of images and depth as input
        batch["image"].shape : batch_size, c, h, w
        batch["depth"].shape : batch_size, 1, h, w

        Assumes all images in a batch are from the same dataset
        
ControlNet-v1-1-nightly/annotator/zoe/zoedepth/trainers/zoedepth_nk_trainer.py