The way that you get D and D_m seems computation heavy.

[Naruto-Sasuke]

def recursive_img(label,res): #Resulution may refers to the final image output i.e. 256x512 or 512x1024
     dim=512 if res>=128 else 1024
#    #M_low will start from 4x8 to resx2*res
     if res == 4:
         downsampled = label #torch.unsqueeze(torch.from_numpy(label).float().permute(2,0,1), dim=0)
     else:
         max1=nn.AvgPool2d(kernel_size=2, padding=0, stride=2)
         downsampled=max1(label)
         img = recursive_img(downsampled, res//2)
         
     global D
     global count
     global D_m

     D.insert(count, downsampled)
     D_m.insert(count, dim)
     count+=1
     return downsampled  

Why not directly assign each D_i and D_m_i with specific values.

[Naruto-Sasuke]

I misread you code.