Fix typo in tests; caught on GPU and TPU (jax-ml#2902)

BuddenD · Apr 30, 2020 · 8d4b685 · 8d4b685
1 parent b39da1f
commit 8d4b685
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/tests/multibackend_test.py b/tests/multibackend_test.py
@@ -165,11 +165,11 @@ def my_sin(x): return np.sin(x)
 
     # jit with `device` spec places the data on the specified device
     result3 = api.jit(my_sin, device=cpus[0])(2)
-    self.assertEqual(result1.device_buffer.device(), cpus[0])
+    self.assertEqual(result3.device_buffer.device(), cpus[0])
 
     # jit with `backend` spec places the data on the specified backend
     result4 = api.jit(my_sin, backend="cpu")(2)
-    self.assertEqual(result1.device_buffer.device(), cpus[0])
+    self.assertEqual(result4.device_buffer.device(), cpus[0])
 
 
 if __name__ == "__main__":