scripts/vsmlrt.py: add tf32 flag to the ort_cuda backend

AmusementClub · Apr 20, 2024 · e1826de · e1826de
1 parent b83941d
commit e1826de
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/scripts/vsmlrt.py b/scripts/vsmlrt.py
@@ -1,4 +1,4 @@
-__version__ = "3.20.8"
+__version__ = "3.20.9"
 
 __all__ = [
     "Backend", "BackendV2",
@@ -101,6 +101,7 @@ class ORT_CUDA:
         fp16_blacklist_ops: typing.Optional[typing.Sequence[str]] = None
         prefer_nhwc: bool = False
         output_format: int = 0 # 0: fp32, 1: fp16
+        tf32: bool = False
 
         # internal backend attributes
         supports_onnx_serialization: bool = True
@@ -2057,6 +2058,7 @@ def _inference(
         if version >= (1, 18, 0):
             kwargs["prefer_nhwc"] = backend.prefer_nhwc
             kwargs["output_format"] = backend.output_format
+            kwargs["tf32"] = backend.tf32
 
         clip = core.ort.Model(
             clips, network_path,