A small update to to ensure ds_loss returns a const (Project-MONAI#5350)

A small fix to followup on Project-MONAI#5338 to ensure ds_loss returns a constant, not an array ### Types of changes <!--- Put an `x` in all the boxes that apply, and remove the not applicable items --> - [x] Non-breaking change (fix or new feature that would not break existing functionality). - [ ] Breaking change (fix or new feature that would cause existing functionality to change). - [ ] New tests added to cover the changes. - [x] Integration tests passed locally by running `./runtests.sh -f -u --net --coverage`. - [x] Quick tests passed locally by running `./runtests.sh --quick --unittests --disttests`. - [ ] In-line docstrings updated. - [ ] Documentation updated, tested `make html` command in the `docs/` folder. Signed-off-by: myron <amyronenko@nvidia.com>

Author: myron

monai/losses/ds_loss.py

@@ -42,21 +42,23 @@ def __init__(self, loss: _Loss, weight_mode: str = "exp", weights: Optional[List
  self.weights = weights
  self.interp_mode = "nearest-exact" if pytorch_after(1, 11) else "nearest"
 
- def get_weight(self, level: int = 0) -> float:
+ def get_weights(self, levels: int = 1) -> List[float]:
  """
- Calculates a weight constant for a given image scale level
+ Calculates weights for a given number of scale levels
  """
- weight = 1.0
- if self.weights is not None and len(self.weights) > level:
- weight = self.weights[level]
+ levels = max(1, levels)
+ if self.weights is not None and len(self.weights) >= levels:
+ weights = self.weights[:levels]
  elif self.weight_mode == "same":
- weight = 1.0
+ weights = [1.0] * levels
  elif self.weight_mode == "exp":
- weight = max(0.5**level, 0.0625)
+ weights = [max(0.5**l, 0.0625) for l in range(levels)]
  elif self.weight_mode == "two":
- weight = 1.0 if level == 0 else 0.5
+ weights = [1.0 if l == 0 else 0.5 for l in range(levels)]
+ else:
+ weights = [1.0] * levels
 
- return weight
+ return weights
 
  def get_loss(self, input: torch.Tensor, target: torch.Tensor):
  """
@@ -71,10 +73,12 @@ def get_loss(self, input: torch.Tensor, target: torch.Tensor):
  def forward(self, input: Union[torch.Tensor, List[torch.Tensor]], target: torch.Tensor):
 
  if isinstance(input, (list, tuple)):
- loss = torch.zeros(1, dtype=torch.float, device=target.device)
+ weights = self.get_weights(levels=len(input))
+ loss = torch.tensor(0, dtype=torch.float, device=target.device)
  for l in range(len(input)):
- loss += self.get_loss(input[l].float(), target) * self.get_weight(l)
+ loss += weights[l] * self.get_loss(input[l].float(), target)
  return loss
+
  return self.loss(input.float(), target)