Update test_batch_states.py

test_gpu() now correctly uses a GPU.
pfnet · Mar 3, 2021 · b5f6907 · b5f6907
1 parent 671b6d6
commit b5f6907
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/tests/utils_tests/test_batch_states.py b/tests/utils_tests/test_batch_states.py
@@ -27,17 +27,17 @@ def phi(state):
         self.assertIsInstance(batch, tuple)
         batch_a, batch_b, batch_c = batch
         np.testing.assert_allclose(
-            batch_a,
+            batch_a.cpu(),
             np.asarray(
                 [
                     [[0, 2], [4, 6]],
                     [[2, 4], [6, 8]],
                 ]
             ),
         )
-        np.testing.assert_allclose(batch_b, np.asarray([0, 1]))
+        np.testing.assert_allclose(batch_b.cpu(), np.asarray([0, 1]))
         np.testing.assert_allclose(
-            batch_c,
+            batch_c.cpu(),
             np.asarray(
                 [
                     [0],