From 8990f1b9aefd89a5096e59b284252201be058e85 Mon Sep 17 00:00:00 2001 From: Yukio Siraichi Date: Mon, 20 May 2024 13:23:05 -0300 Subject: [PATCH] [benchmarks] Add default value to `move_to_device`. (#7080) --- benchmarks/util.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/benchmarks/util.py b/benchmarks/util.py index 88bf452bbdd..8ab1d5a0181 100644 --- a/benchmarks/util.py +++ b/benchmarks/util.py @@ -76,7 +76,7 @@ def is_xla_device_available(devkind): return r.returncode == 0 -def move_to_device(item, device, torch_xla2): +def move_to_device(item, device, torch_xla2=False): if torch_xla2: import torch_xla2 import jax