admin/lerobot
1
0
Fork 0
mirror of https://github.com/huggingface/lerobot.git synced 2026-07-23 17:56:07 +00:00
Code Issues Packages Projects Releases Wiki Activity

cast float64 to float32 for mps

Browse Source
This commit is contained in:
Pepijn
2025-09-16 16:48:08 +02:00
parent 5924d4d9eb
commit 8f624f1c1e
1 changed files with 2 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/lerobot/policies/pi05_openpi/modeling_pi05openpi.py
+2
View File
@@ -37,6 +37,8 @@ from lerobot.policies.pretrained import PreTrainedPolicy
# Helper functions # Helper functions
def get_safe_dtype(target_dtype, device_type): # see openpi `get_safe_dtype` (exact copy) def get_safe_dtype(target_dtype, device_type): # see openpi `get_safe_dtype` (exact copy)
"""Get a safe dtype for the given device type.""" """Get a safe dtype for the given device type."""
if device_type == "mps" and target_dtype == torch.float64:
return torch.float32
if device_type == "cpu": if device_type == "cpu":
# CPU doesn't support bfloat16, use float32 instead # CPU doesn't support bfloat16, use float32 instead
if target_dtype == torch.bfloat16: if target_dtype == torch.bfloat16: