diff --git a/monai/networks/nets/segresnet_ds.py b/monai/networks/nets/segresnet_ds.py index 1ac5a79ee3..098e490511 100644 --- a/monai/networks/nets/segresnet_ds.py +++ b/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) diff --git a/monai/networks/nets/vista3d.py b/monai/networks/nets/vista3d.py index 4215a9a594..6313b7812d 100644 --- a/monai/networks/nets/vista3d.py +++ b/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):