Add default value to move_to_device.

pytorch · May 18, 2024 · 62c3ba6 · 62c3ba6
1 parent 3c59087
commit 62c3ba6
Showing 1 changed file with 1 addition and 1 deletion.
diff --git 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