Fixed failing torch utils test.

Unity-Technologies · Aug 25, 2024 · 6ad2bdf · 6ad2bdf
1 parent fac2c2c
commit 6ad2bdf
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/ml-agents/mlagents/trainers/tests/test_torch_utils.py b/ml-agents/mlagents/trainers/tests/test_torch_utils.py
@@ -10,13 +10,13 @@
 @pytest.mark.parametrize(
     "device_str, expected_type, expected_index, expected_tensor_type",
     [
-        ("cpu", "cpu", None, torch.FloatTensor),
+        ("cpu", "cpu", None, torch.float32),
         ("cuda", "cuda", None, torch.cuda.FloatTensor),
         ("cuda:42", "cuda", 42, torch.cuda.FloatTensor),
-        ("opengl", "opengl", None, torch.FloatTensor),
+        ("opengl", "opengl", None, torch.float32),
     ],
 )
-@mock.patch.object(torch, "set_default_tensor_type")
+@mock.patch.object(torch, "set_default_dtype")
 def test_set_torch_device(
     mock_set_default_tensor_type,
     device_str,