add metadata of device and cuda version

microsoft · Nov 28, 2023 · 15476ba · 15476ba
1 parent 6741855
commit 15476ba
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/onnxruntime/python/tools/transformers/models/stable_diffusion/demo_utils.py b/onnxruntime/python/tools/transformers/models/stable_diffusion/demo_utils.py
@@ -326,6 +326,8 @@ def get_metadata(args, is_xl: bool = False) -> Dict[str, Any]:
         except PackageNotFoundError:
             continue
     metadata["packages"] = packages
+    metadata["device"] = torch.cuda.get_device_name()
+    metadata["torch.version.cuda"] = torch.version.cuda
 
     return metadata