diff --git a/benchmarks/torchbench_model.py b/benchmarks/torchbench_model.py index d715ee6dab0..e60e3b21ed5 100644 --- a/benchmarks/torchbench_model.py +++ b/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