remove broken energy tracker reset method

huggingface · Oct 14, 2024 · 9fdb506 · 9fdb506
1 parent b7b8dbb
commit 9fdb506
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 3 deletions.
diff --git a/examples/pytorch_llama.py → examples/pytorch_gpt2.py b/examples/pytorch_llama.py → examples/pytorch_gpt2.py
diff --git a/examples/pytorch_llama.yaml → examples/pytorch_gpt2.yaml b/examples/pytorch_llama.yaml → examples/pytorch_gpt2.yaml
@@ -15,6 +15,7 @@ launcher:
 backend:
   model: gpt2
   device: cuda
+  device_ids: 0
   torch_dtype: float16
 
 scenario:

diff --git a/tests/test_examples.py b/tests/test_examples.py
@@ -48,7 +48,7 @@
 # can be run with pytest tests/test_example.py -s -k "cuda and pytorch"
 CUDA_PYTORCH_CONFIGS = [
     "pytorch_bert.yaml",
-    "pytorch_llama.yaml",
+    "pytorch_gpt2.yaml",
 ]
 
 # can be run with pytest tests/test_example.py -s -k "cuda and txi"
@@ -58,7 +58,7 @@
 
 # can be run with pytest tests/test_example.py -s -k "cuda and tensorrt_llm"
 CUDA_TENSORRT_LLM_CONFIGS = [
-    "trt_llama.yaml",
+    # "trt_llama.yaml",
 ]
 
 # can be run with pytest tests/test_example.py -s -k "cuda and vllm"
@@ -72,14 +72,16 @@
     "llama_cpp_text_generation.yaml",
 ]
 
+# can be run with pytest tests/test_example.py -s -k "cuda and pytorch"
 CUDA_PYTORCH_SCRIPTS = [
     "pytorch_bert.py",
-    "pytorch_llama.py",
+    "pytorch_gpt2.py",
 ]
 
 # Those tests are not run on the CI/CD pipeline as they are currently broken
 UNTESTED_YAML_CONFIGS = [
     "energy_star.yaml",
+    "trt_llama.yaml",
 ]
 
 ALL_YAML_CONFIGS = (