fix

mlcommons · Oct 10, 2023 · fea7529 · fea7529
1 parent 7ba0c4c
commit fea7529
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 7 deletions.
diff --git a/algorithmic_efficiency/logger_utils.py b/algorithmic_efficiency/logger_utils.py
@@ -14,11 +14,11 @@
 from algorithmic_efficiency import spec
 from algorithmic_efficiency.pytorch_utils import pytorch_setup
 
-try:
-  import wandb  # pylint: disable=g-import-not-at-top
-except ModuleNotFoundError:
-  logging.exception('Unable to import wandb.')
-  wandb = None
+# try:
+#   import wandb  # pylint: disable=g-import-not-at-top
+# except ModuleNotFoundError:
+#   logging.exception('Unable to import wandb.')
+wandb = None
 
 USE_PYTORCH_DDP, RANK, DEVICE, N_GPUS = pytorch_setup()
 

diff --git a/algorithmic_efficiency/workloads/fastmri/fastmri_jax/workload.py b/algorithmic_efficiency/workloads/fastmri/fastmri_jax/workload.py
@@ -78,8 +78,8 @@ def loss_fn(self,
               label_smoothing: float = 0.0) -> spec.Tensor:  # differentiable
     del label_smoothing
     losses = jnp.sum(
-        jnp.abs(outputs_batch - label_batch),
-        axis=tuple(range(1, outputs_batch.ndim)))
+        jnp.abs(logits_batch - label_batch),
+        axis=tuple(range(1, logits_batch.ndim)))
     # mask_batch is assumed to be shape [batch].
     if mask_batch is not None:
       losses *= mask_batch