ttrt add --seed argument for rand tensor data (#295)

tenstorrent · Aug 5, 2024 · 5a99c93 · 5a99c93
1 parent f818088
commit 5a99c93
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 0 deletions.
diff --git a/runtime/tools/python/ttrt/__init__.py b/runtime/tools/python/ttrt/__init__.py
@@ -80,6 +80,11 @@ def main():
         action="store_true",
         help="save all artifacts during run",
     )
+    run_parser.add_argument(
+        "--seed",
+        default=0,
+        help="Seed for random number generator",
+    )
     run_parser.add_argument("binary", help="flatbuffer binary file")
     run_parser.set_defaults(func=run)
 

diff --git a/runtime/tools/python/ttrt/common/api.py b/runtime/tools/python/ttrt/common/api.py
@@ -133,6 +133,8 @@ def run(args):
     device = ttrt.runtime.open_device(device_ids)
     atexit.register(lambda: ttrt.runtime.close_device(device))
 
+    torch.manual_seed(args.seed)
+
     for (binary_name, fbb, fbb_dict) in fbb_list:
         torch_inputs[binary_name] = []
         torch_outputs[binary_name] = []