Update ssl_head.py (Project-MONAI#91)

Model name needs to be changed at line 88 to match the change made in the last update at line 23.
KumoLiu · Aug 13, 2022 · 32c8209 · 32c8209
1 parent 659d711
commit 32c8209
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/SwinUNETR/Pretrain/models/ssl_head.py b/SwinUNETR/Pretrain/models/ssl_head.py
@@ -85,7 +85,7 @@ def __init__(self, args, upsample='vae', dim=768):
                                       )
 
     def forward(self, x):
-        x_out = self.swin_vit(x.contiguous())[4]
+        x_out = self.swinViT(x.contiguous())[4]
         _, c, h, w, d = x_out.shape
         x4_reshape = x_out.flatten(start_dim=2, end_dim=4)
         x4_reshape = x4_reshape.transpose(1, 2)