Minor fix

NKI-AI · Oct 20, 2024 · a96907a · a96907a
1 parent 930e2e5
commit a96907a
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/direct/nn/transformers/mri_transformers.py b/direct/nn/transformers/mri_transformers.py
@@ -625,7 +625,7 @@ def __init__(
             Whether to compute the output per coil.
         """
         super().__init__()
-        self.transformer = VisionTransformer2D(
+        self.transformer = VisionTransformer3D(
             average_img_size=average_size,
             patch_size=patch_size,
             in_channels=COMPLEX_SIZE,