Changed default cuda dtype to torch.float32.

Unity-Technologies · Oct 4, 2024 · 54664fc · 54664fc
1 parent c683e77
commit 54664fc
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/ml-agents/mlagents/torch_utils/torch.py b/ml-agents/mlagents/torch_utils/torch.py
@@ -53,7 +53,7 @@ def set_torch_config(torch_settings: TorchSettings) -> None:
 
     if _device.type == "cuda":
         torch.set_default_device(_device.type)
-        torch.set_default_dtype(torch.cuda.FloatTensor)
+        torch.set_default_dtype(torch.float32)
     else:
         torch.set_default_dtype(torch.float32)
     logger.debug(f"default Torch device: {_device}")