add variants to workload registry

mlcommons · Dec 5, 2023 · 145d649 · 145d649
1 parent 4cabfa6
commit 145d649
Showing 1 changed file with 12 additions and 0 deletions.
diff --git a/algorithmic_efficiency/workloads/workloads.py b/algorithmic_efficiency/workloads/workloads.py
@@ -36,6 +36,18 @@
         'workload_path': 'fastmri/fastmri',
         'workload_class_name': 'FastMRIWorkload',
     },
+    'fastmri_model_size': {
+        'workload_path': 'fastmri/fastmri',
+        'workload_class_name': 'FastMRIModelSizeWorkload',
+    },
+    'fastmri_tanh': {
+        'workload_path': 'fastmri/fastmri',
+        'workload_class_name': 'FastMRITanhWorkload',
+    },
+    'fastmri_layernorm': {
+        'workload_path': 'fastmri/fastmri',
+        'workload_class_name': 'FastMRILayerNormWorkload',
+    },
     'imagenet_resnet': {
         'workload_path': 'imagenet_resnet/imagenet',
         'workload_class_name': 'ImagenetResNetWorkload',