Question about memory consumption

[Wxsxulin]

Thank you for your excellent work. I have a question about the memory consumption of demotion blur。I used two A100s（40G）to train the model of demotion blur，but the memory consumption can‘t support the batch_size = 1,patch_size = 512. I didn't change any parameters in the code，could you help me to find the possible problems？