diff --git a/turbodiffusion/rcm/datasets/utils.py b/turbodiffusion/rcm/datasets/utils.py
index e5b3672..2a7561b 100644
--- a/turbodiffusion/rcm/datasets/utils.py
+++ b/turbodiffusion/rcm/datasets/utils.py
@@ -14,7 +14,7 @@
 # limitations under the License.
 
 # "width:height"
-IMAGE_RES_SIZE_INFO: dict[str, tuple[int, int]] = {
+IMAGE_RES_SIZE_INFO: dict[str, dict[str, tuple[int, int]]] = {
     "1024": {"1:1": (1024, 1024), "4:3": (1168, 880), "3:4": (880, 1168), "16:9": (1360, 768), "9:16": (768, 1360)},
     "720": {"1:1": (960, 960), "4:3": (960, 704), "3:4": (704, 960), "16:9": (1280, 704), "9:16": (704, 1280)},
     "512": {"1:1": (512, 512), "4:3": (640, 512), "3:4": (512, 640), "16:9": (640, 384), "9:16": (384, 640)},
@@ -22,7 +22,7 @@
 }
 
 # "width:height"
-VIDEO_RES_SIZE_INFO: dict[str, tuple[int, int]] = {
+VIDEO_RES_SIZE_INFO: dict[str, dict[str, tuple[int, int]]] = {
     "720": {"1:1": (960, 960), "4:3": (960, 704), "3:4": (704, 960), "16:9": (1280, 704), "9:16": (704, 1280)},
     "512": {"1:1": (512, 512), "4:3": (640, 512), "3:4": (512, 640), "16:9": (640, 384), "9:16": (384, 640)},
     "480": {"1:1": (480, 480), "4:3": (640, 480), "3:4": (480, 640), "16:9": (768, 432), "9:16": (432, 768)},