confirmed
w/ a bunch of print statements in the official model here:
@@ -310,7 +310,6 @@ class UNetModel:
]
self.input_blocks = [
[Conv2d(4, 320, kernel_size=3, padding=1)],
- # TODO: my head sizes and counts are a guess
[ResBlock(320, 1280, 320), SpatialTransformer(320, 768, 8, 40)],
[ResBlock(320, 1280, 320), SpatialTransformer(320, 768, 8, 40)],
[Downsample(320)],
