Sagemaker Jumpstar, Stable Diffusion x4 upscaler FP16, OOM error

Hello,

I've deployed Stable Diffusion x4 upscaler FP16 model via Sagemaker. Using ml.g5.2xlarge instance as advised by the tutorial (Introduction to JumpStart - Enhance image quality guided by prompt). When using 512 x 512 image I get CUDA Out of memory error. Is this the expected behaviour because 512 x 512 is too large for this model or it should be working fine and there's something else that's causing OOM error?

Has anyone used this model to successfully upscale the 512 x 512 or larger images?

Best regards, Victor