Update metrics.py

The problem with torchrun and jax seems to be caused by jax.nn.sigmoid.

Author: davidtweedle

algoperf/workloads/ogbg/metrics.py

@@ -37,21 +37,23 @@ def compute(self):
     labels = values['labels']
     logits = values['logits']
     mask = values['mask']
+    sigmoid = jax.nn.sigmoid
 
     if USE_PYTORCH_DDP:
       # Sync labels, logits, and masks across devices.
-      all_values = [np.array(labels), np.array(logits), np.array(mask)]
+      all_values = [labels, logits, mask]
       for idx, array in enumerate(all_values):
         tensor = torch.as_tensor(array, device=DEVICE)
         # Assumes that the tensors on all devices have the same shape.
         all_tensors = [torch.zeros_like(tensor) for _ in range(N_GPUS)]
         dist.all_gather(all_tensors, tensor)
         all_values[idx] = torch.cat(all_tensors).cpu().numpy()
       labels, logits, mask = all_values
+      sigmoid = lambda x: 1 / (1 + np.exp(-x))
 
     mask = mask.astype(bool)
 
-    probs = 1 / (1 + np.exp(-logits))
+    probs = sigmoid(logits)
     num_tasks = labels.shape[1]
     average_precisions = np.full(num_tasks, np.nan)