Optimize VISTA3D

monai/networks/nets/segresnet_ds.py

@@ -508,8 +508,10 @@ def forward(  # type: ignore
 
         outputs: list[torch.Tensor] = []
         outputs_auto: list[torch.Tensor] = []
-        x_ = x.clone()
+        x_ = x
         if with_point:
+            if with_label:
+                x_ = x.clone()
             i = 0
             for level in self.up_layers:
                 x = level["upsample"](x)

monai/networks/nets/vista3d.py

@@ -639,12 +639,10 @@ def forward(self, src: torch.Tensor, class_vector: torch.Tensor):
         if self.use_mlp:
             class_embedding = self.mlp(class_embedding)
         # [b,1,feat] @ [1,feat,dim], batch dimension become class_embedding batch dimension.
-        masks = []
-        for i in range(b):
-            mask = class_embedding @ src[[i]].view(1, c, h * w * d)
-            masks.append(mask.view(-1, 1, h, w, d))
+        masks_embedding = class_embedding.squeeze() @ src.view(b, c, h * w * d)
+        masks_embedding = masks_embedding.view(b, -1, h, w, d).transpose(0, 1)
 
-        return torch.cat(masks, 1), class_embedding
+        return masks_embedding, class_embedding
 
 
 class TwoWayTransformer(nn.Module):