From 145d649287a6cba5ee9eb420cff0b7ad32571fc7 Mon Sep 17 00:00:00 2001
From: Priya Kasimbeg <kasimbeg@google.com>
Date: Tue, 5 Dec 2023 04:41:58 +0000
Subject: [PATCH] add variants to workload registry

---
 algorithmic_efficiency/workloads/workloads.py | 12 ++++++++++++
 1 file changed, 12 insertions(+)

diff --git a/algorithmic_efficiency/workloads/workloads.py b/algorithmic_efficiency/workloads/workloads.py
index 679cc75d7..5631841d5 100644
--- 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',