def forward(self, details, norms, max_pixels, min_pixels):

    return torch.FloatTensor([[torch.where(dtl_pixel < 0, weight(dtl_pixel, nrm, max_pixel, min_pixel), dtl_pixel) for j, (dtl_pixel, nrm, max_pixel, min_pixel) in
                               enumerate(zip(detail, norm, max_pixel, min_pixel))] for i, (detail, norm, max_pixel, min_pixel) in
                              enumerate(zip(details, norms, max_pixels, min_pixels))]).reshape((1024, 1024))

I organized dataset (with torch.utils.data.Dataset) for details, norms, min pixels, max pixels as separated class, that I calculate in the beginning, and is used by the kernel.
weight function is:
weight = lambda detail, norm, max, min: detail / torch.exp(torch.pow((max * self.lamda) / (norm * min + self.epsilon), 0.01))