cleanup

2026-09-19 09:14:35 +02:00 · 2026-01-10 20:48:36 +03:00
parent f4d87d6cdf
commit fa2917a6be
1 changed files with 1 additions and 1 deletions
@@ -70,7 +70,7 @@ original_get_autocast_dtype = torch.get_autocast_dtype
@wraps(torch.get_autocast_dtype)
 def torch_get_autocast_dtype(device_type=None):
    if device_type is None or check_cuda(device_type) or check_device_type(device_type, "xpu"):
-        return torch.bfloat16
+        return devices.dtype or torch.bfloat16
    else:
        return original_get_autocast_dtype(device_type)