From 4af702e8ad19690886ccda0c6785df95b66f68aa Mon Sep 17 00:00:00 2001 From: priyakasimbeg Date: Thu, 23 Nov 2023 00:15:56 +0000 Subject: [PATCH] torch compile --- submission_runner.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/submission_runner.py b/submission_runner.py index 9928f9b84..d4d9e4ce7 100644 --- a/submission_runner.py +++ b/submission_runner.py @@ -221,9 +221,12 @@ def train_once( model_init_rng, dropout_rate, aux_dropout_rate) if FLAGS.framework == 'pytorch' and FLAGS.torch_compile: compile_error_workloads = [ - 'librispeech_conformer', 'librispeech_conformer_gelu', + 'librispeech_conformer', + 'librispeech_conformer_gelu', + 'librispeech_conformer_layer_norm', + 'librispeech_conformer_attention_temperature', 'ogbg', 'criteo1tb', - 'imagenet_vit' + 'imagenet_vit', ] eager_backend_workloads = ['librispeech_deepspeech'] aot_eager_backend_workloads = []