Updated requirements, batch size.

pytorch · Jul 14, 2023 · ba57c50 · ba57c50
1 parent 89f4497
commit ba57c50
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 5 deletions.
diff --git a/torchbenchmark/models/hf_Whisper/__init__.py b/torchbenchmark/models/hf_Whisper/__init__.py
@@ -4,9 +4,9 @@
 
 class Model(HuggingFaceModel):
     task = SPEECH.RECOGNITION
-    # https://cdn.openai.com/papers/whisper.pdf Says for large-v2 they trained on 1024 batch sizes. 
-    DEFAULT_TRAIN_BSIZE = 8
-    DEFAULT_EVAL_BSIZE = 8
+    # https://cdn.openai.com/papers/whisper.pdf Says for large-v2 they trained on 1024 batch sizes, with 16 GPUs
+    DEFAULT_TRAIN_BSIZE = 64
+    DEFAULT_EVAL_BSIZE = 64
 
     def __init__(self, test, device, jit=False, batch_size=None, extra_args=[]):
         super().__init__(name="hf_Whisper", test=test, device=device, jit=jit, batch_size=batch_size, extra_args=extra_args)

diff --git a/torchbenchmark/models/hf_Whisper/requirements.txt b/torchbenchmark/models/hf_Whisper/requirements.txt
@@ -1,2 +1 @@
-sentencepiece
-datasets
+numba