Use BENCHMARK_USE_SGD set, instead.

benchmarks/torchbench_model.py

@@ -32,22 +32,30 @@
 # torchbench models that might OOM using Adam.
 # This list was extracted from PyTorch's repository: benchmarks/dynamo/common.py
 TRAIN_WITH_SGD = {
-    "torchrec_dlrm",
+    "BERT_pytorch",
+    "LearningToPaint",
+    "alexnet",
+    "dcgan",
     "demucs",
-    "hf_T5_base",
-    "hf_clip",
-    "llama_v2_7b_16h",
-    "mobilenet_v2_quantized_qat",
-    "phi_1_5",
-    "resnet50_quantized_qat",
-    "DALLE2_pytorch",
-    "moco",
-    "timm_efficientdet",
-    "pytorch_CycleGAN_and_pix2pix",
-    "vision_maskrcnn",
+    "densenet121",
     "dlrm",
-    "resnet50",
-    *DETECTRON2_MODELS,
+    "fastNLP_Bert",
+    "mobilenet_v2",
+    "phlippe_densenet",
+    "phlippe_resnet",
+    "pytorch_stargan",
+    "resnet18",
+    "shufflenet_v2_x1_0",
+    "speech_transformer",
+    "squeezenet1_1",
+    "stable_diffusion_text_encoder",
+    "timm_efficientdet",
+    "timm_nfnet",
+    "timm_regnet",
+    "timm_vision_transformer",
+    "timm_vovnet",
+    "vgg16",
+    "hf_T5",
 }
 
 # Skip the experiment of a model if any of the experiment configs in the list is fully matched