fix: only cast float64 tensors in dict_to_device

Avoid converting boolean masks to float, which breaks ~ and torch.where.

Co-authored-by: Cursor <cursoragent@cursor.com>

Author: georgeyiasemis

direct/utils/__init__.py

@@ -172,7 +172,13 @@ def dict_to_device(
     """
     if keys is None:
         keys = data.keys()
-    return {k: v.float().to(device) if isinstance(v, torch.Tensor) else v for k, v in data.items() if k in keys}
+    return {
+        k: (v.float() if v.dtype == torch.float64 else v).to(device)
+        if isinstance(v, torch.Tensor)
+        else v
+        for k, v in data.items()
+        if k in keys
+    }
 
 
 def detach_dict(data: Dict[str, torch.Tensor], keys: Optional[Union[List, Tuple, KeysView]] = None) -> Dict: