Merge pull request #38 from arjunsuresh/mlperf-inference

Support cuda in docker container for reference implementations
GATEOverflow · May 31, 2024 · c48f383 · c48f383
2 parents a341ccc + b285782
commit c48f383
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/script/app-mlperf-inference/_cm.yaml b/script/app-mlperf-inference/_cm.yaml
@@ -900,6 +900,10 @@ variations:
     add_deps_recursive:
       mlperf-inference-implementation:
         tags: _cpu
+
+  cuda,reference:
+    docker:
+      base_image: nvcr.io/nvidia/pytorch:24.03-py3
   cuda:
     docker:
       all_gpus: 'yes'