fix masked_mean() error expected x and y to be on same device

arises due to torch.tensor(float("nan")) defaulting to CPU

Author: janosh

aviary/core.py

@@ -608,14 +608,11 @@ def masked_mean(x: torch.Tensor, mask: torch.BoolTensor, dim: int = 0) -> torch.
         x (torch.Tensor): Tensor to compute standard deviation of.
         mask (torch.BoolTensor): Same shape as x with True where x is valid and False
             where x should be masked. Mask should not be all False in any column of
-            dimension dim to avoid NaNs.
+            dimension dim to avoid NaNs from zero division.
         dim (int, optional): Dimension to take mean of. Defaults to 0.
 
     Returns:
         torch.Tensor: Same shape as x, except dimension dim reduced.
     """
-    # mask should be True where x is valid and False where x should be masked
-    x_nan = torch.where(mask, x, torch.tensor(float("nan")))
-    # torch.tensor(float("nan")) can be simplified to torch.nan in torch>=1.12
-
+    x_nan = x.masked_fill(~mask, float("nan"))
     return x_nan.nanmean(dim=dim)

examples/mat_bench/slurm_submit.py

@@ -43,9 +43,7 @@
 print(f"{{job_id=}}")
 print("{model_name=}")
 
-job_array_id = os.environ.get("SLURM_ARRAY_TASK_ID")
-if job_array_id is not None:
-    job_array_id = int(job_array_id)
+job_array_id = int(os.environ.get("SLURM_ARRAY_TASK_ID"), 0)
 print(f"{{job_array_id=}}")
 
 dataset_name, fold = list(product({datasets}, {folds}))[job_array_id]

examples/mp_wbm/slurm_submit.py

@@ -39,9 +39,7 @@
 print("{model_name=}")
 print("{data_path=}")
 
-job_array_id = os.environ.get("SLURM_ARRAY_TASK_ID")
-if job_array_id is not None:
-    job_array_id = int(job_array_id)
+job_array_id = int(os.environ.get("SLURM_ARRAY_TASK_ID"), 0)
 print(f"{{job_array_id=}}")
 
 run_wrenformer_on_mp_wbm(