diff --git a/pathwaysutils/profiling.py b/pathwaysutils/profiling.py index 7e067bc..029fc83 100644 --- a/pathwaysutils/profiling.py +++ b/pathwaysutils/profiling.py @@ -141,7 +141,7 @@ def _create_profile_request( if isinstance(v, bool): advanced_config[k] = {"boolValue": v} elif isinstance(v, int): - advanced_config[k] = {"intValue": v} + advanced_config[k] = {"int64Value": v} elif isinstance(v, str): advanced_config[k] = {"stringValue": v} else: diff --git a/pathwaysutils/test/profiling_test.py b/pathwaysutils/test/profiling_test.py index c189fb0..d6d926b 100644 --- a/pathwaysutils/test/profiling_test.py +++ b/pathwaysutils/test/profiling_test.py @@ -569,6 +569,8 @@ def test_create_profile_request_with_options(self): "tpu_num_chips_to_profile_per_task": 3, "tpu_num_sparse_core_tiles_to_trace": 5, "tpu_trace_mode": "TRACE_COMPUTE", + "tpu_num_sparse_cores_to_trace": 1, + "tpu_enable_flag": True, } request = profiling._create_profile_request( @@ -586,9 +588,15 @@ def test_create_profile_request_with_options(self): "pwTraceOptions": { "enablePythonTracer": True, "advancedConfiguration": { - "tpu_num_chips_to_profile_per_task": {"intValue": 3}, - "tpu_num_sparse_core_tiles_to_trace": {"intValue": 5}, + "tpu_num_chips_to_profile_per_task": { + "int64Value": 3 + }, + "tpu_num_sparse_core_tiles_to_trace": { + "int64Value": 5 + }, "tpu_trace_mode": {"stringValue": "TRACE_COMPUTE"}, + "tpu_num_sparse_cores_to_trace": {"int64Value": 1}, + "tpu_enable_flag": {"boolValue": True}, }, }, },