amend yaml config example

NVIDIA · Jun 18, 2024 · 1fca119 · 1fca119
1 parent 62f91a9
commit 1fca119
Showing 1 changed file with 11 additions and 6 deletions.
diff --git a/tests/test_config.py b/tests/test_config.py
@@ -19,31 +19,36 @@
 plugins:
   generators:
     huggingface:
-      dtype: general
-      gpu: 0
+      hf_args:
+        torch_dtype: float16
       Pipeline:
-        dtype: bfloat16
+        hf_args:
+            device: cuda
   probes:
     test:
       generators:
         huggingface:
             Pipeline:
-              dtype: for_probe
+                hf_args:
+                    torch_dtype: float16
   detector:
       test:
         val: tests
         Blank:
           generators:
             huggingface:
-                gpu: 1
+                hf_args:
+                    torch_dtype: float16
+                    device: cuda:1
                 Pipeline:
                   dtype: for_detector
   buffs:
       test:
         Blank:
           generators:
             huggingface:
-                gpu: 1
+                hf_args:
+                    device: cuda:0
                 Pipeline:
                   dtype: for_detector
 """.encode(