mahiratmis
commented
3 months ago should I call restoration method like this?
y_t = torch.randn_like(self.gt_image) # y_t will be of shape B,3,H,W
self.output, self.visuals = self.netG.restoration(self.cond_image, y_t, sample_num=self.sample_num)
Hello thank you for the code and especially for the distributed support it contains. I try to enlight dark images. In my problem settings the y_cond is of shape B, 32, H, W and the y_0 (ground truth) is of shape B, 3, H, W. Therefore, I set the in_channel value of the Unet to 35. The training went smoothly. However, during testing inside restoration method the code is as follows: y_t = default(y_t, lambda: torch.randn_like(y_cond)) Since y_t is None at the beginning, the code above creates y_t as a random noise of shape B, 32, H, W and inside p_sample y_t is concatenated with y_cond which results in a shape of B, 64, H, W. That sahpe is not compatible with the expected shape of the input to the Unet.
What am I missing? Can someone please help ?