Fixed failing GPU test.

Unity-Technologies · Oct 4, 2024 · ee94d91 · ee94d91
1 parent 9d34978
commit ee94d91
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/ml-agents/mlagents/trainers/tests/torch_entities/test_attention.py b/ml-agents/mlagents/trainers/tests/torch_entities/test_attention.py
@@ -217,7 +217,7 @@ def test_predict_minimum_training():
             argmin = argmin.squeeze()
             argmin = argmin.detach()
         sliced_oh = onehots[:, : num + 1]
-        inp = torch.cat([inp, sliced_oh], dim=2)
+        inp = torch.cat([inp, sliced_oh.to(torch.get_default_device())], dim=2)
 
         embeddings = entity_embedding(inp, inp)
         masks = get_zero_entities_mask([inp])