fix tests about nan_targets

graphium/trainer/losses.py

@@ -62,8 +62,9 @@ def forward(self, input: Tensor, target: Tensor, nan_targets: Tensor = None) ->
         target = target.flatten()
 
         # set input and target with nans to 0s for regression loss
-        input[nan_targets] = 0.0
-        target[nan_targets] = 0.0
+        if nan_targets is not None:
+            input[nan_targets] = 0.0
+            target[nan_targets] = 0.0
         # regression loss needs normalized logits to probability as input to do inner product with self.brackets
         # we apply softmax on the raw logits first
         softmax_input = self.softmax(input)
@@ -74,14 +75,15 @@ def forward(self, input: Tensor, target: Tensor, nan_targets: Tensor = None) ->
         regression_input = torch.inner(softmax_input, self.brackets.to(input.device))
         regression_loss = self.regression_loss(regression_input, target.float(), reduction=self.reduction)
         # regression_loss needs some scaling by total_targets/num_real_targets
-        num_real_targets = (~nan_targets).sum()
-        factor1 = torch.where(num_real_targets > 0, 1, 0)
-        factor2 = torch.where(num_real_targets > 0, 0, 1)
-        regression_loss = factor1 * regression_loss * nan_targets.numel() / (num_real_targets + factor2)
+        if nan_targets is not None:
+            num_real_targets = (~nan_targets).sum()
+            factor1 = torch.where(num_real_targets > 0, 1, 0)
+            factor2 = torch.where(num_real_targets > 0, 0, 1)
+            regression_loss = factor1 * regression_loss * nan_targets.numel() / (num_real_targets + factor2)
 
-        # set input and target with nans to -1000s for ce loss
-        input[nan_targets] = -1000
-        target[nan_targets] = -1000
+            # set input and target with nans to -1000s for ce loss
+            input[nan_targets] = -1000
+            target[nan_targets] = -1000
         # cross_entropy loss needs raw logits as input
         # ce_loss does not need scaling as it already ignores -1000 masked nan values
         ce_loss = F.cross_entropy(