model diff tests fastmri

mlcommons · Dec 7, 2023 · 8af0c73 · 8af0c73
1 parent 993c8da
commit 8af0c73
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/tests/modeldiffs/fastmri_layernorm/compare.py b/tests/modeldiffs/fastmri_layernorm/compare.py
@@ -8,9 +8,9 @@
 
 from algorithmic_efficiency import spec
 from algorithmic_efficiency.workloads.fastmri.fastmri_jax.workload import \
-    FastMRIWorkload as JaxWorkload
+    FastMRILayerNormWorkload as JaxWorkload
 from algorithmic_efficiency.workloads.fastmri.fastmri_pytorch.workload import \
-    FastMRIWorkload as PytWorkload
+    FastMRILayerNormWorkload as PytWorkload
 from tests.modeldiffs.diff import out_diff
 
 

diff --git a/tests/modeldiffs/fastmri_tanh/compare.py b/tests/modeldiffs/fastmri_tanh/compare.py
@@ -8,9 +8,9 @@
 
 from algorithmic_efficiency import spec
 from algorithmic_efficiency.workloads.fastmri.fastmri_jax.workload import \
-    FastMRIWorkload as JaxWorkload
+    FastMRITanhWorkload as JaxWorkload
 from algorithmic_efficiency.workloads.fastmri.fastmri_pytorch.workload import \
-    FastMRIWorkload as PytWorkload
+    FastMRITanhWorkload as PytWorkload
 from tests.modeldiffs.diff import out_diff