diff --git a/scripts/upload_scribe.py b/scripts/upload_scribe.py index 392adca69..51f11b2dd 100644 --- a/scripts/upload_scribe.py +++ b/scripts/upload_scribe.py @@ -89,6 +89,7 @@ def __init__(self): "git_commit_time", "git_dirty", "pytorch_version", + "torchvision_version", "python_version", "machine_kernel", "machine_processor", diff --git a/scripts/upload_scribe_v2.py b/scripts/upload_scribe_v2.py index 441ebd548..2a0ea0893 100644 --- a/scripts/upload_scribe_v2.py +++ b/scripts/upload_scribe_v2.py @@ -122,6 +122,7 @@ def __init__(self): "git_commit_time", "git_dirty", "pytorch_version", + "torchvision_version", "python_version", "machine_kernel", "machine_processor", diff --git a/test_upload_scribe.py b/test_upload_scribe.py new file mode 100644 index 000000000..1147b1328 --- /dev/null +++ b/test_upload_scribe.py @@ -0,0 +1,11 @@ +import importlib + + +def test_scribe_uploaders_accept_torchvision_version(): + for module_name in ("scripts.upload_scribe", "scripts.upload_scribe_v2"): + module = importlib.import_module(module_name) + message = module.PytorchBenchmarkUploader().format_message( + {"time": 1, "torchvision_version": "0.18.0"} + ) + + assert message["normal"]["torchvision_version"] == "0.18.0"