SSIM loss for standardized data

[RohanKPatil]

It is usually a practice to standardize the data for deeper models for training stability and faster convergence. Its difficult to compute the SSIM score for standardized data, especially in case of medical datasets - since we standardize each slice/sample with respect to the entire volume from which that slice/sample came (mean and std of the volume is used). I have a modified dataloader which can give the images for each batch along with the metadata, which contains the mean and std of the original volume for the corresponding sample, which can be used to denormalize the data and then compute the SSIM score. Is there any cleaner way to do this, given we get the metadata for each sample of the batch in the training_step function?

def training_step(self, batch, batch_idx):
        inputs, targets, meta = batch['input'], batch['target'], batch['meta'] 
        outputs = self(inputs)
        targets = targets.to(dtype=outputs.dtype)
        ssim_batch = []
        for i in range(inputs.size(0)):
            X = (outputs[i].unsqueeze(0) * meta['std'][i]) + meta['mean'][i]
            Y = (targets[i].unsqueeze(0) * meta['std'][i]) + meta['mean'][i]
            inst_ssim = self.criterion_ssim(X, Y)    # inst_ssim is ssim loss for each sample in the batch
            ssim_batch.append(inst_ssim)
        ssim_loss = 1 - torch.Tensor(ssim_batch).mean()
        self.log('train_loss', loss)
        return {'loss': loss}

This way is definitely not recommended, since we are changing tensors. Is there a better way to do this?