pytorch - 💡(How to fix) Fix torch.linspace with integer dtype produces inconsistent results on CPU and CUDA [1 comments, 2 participants]

TL;DR

The discrepancy in torch.linspace results between CPU and CUDA for dtype=torch.int64 may be due to different rounding behaviors, and a potential workaround is to use a consistent rounding method.

Guidance

• Verify that the issue persists across different PyTorch and CUDA versions to rule out version-specific bugs.
• Consider using the torch.round function to explicitly round the results of torch.linspace to the nearest integer, ensuring consistency between CPU and CUDA.
• Check if using dtype=torch.float64 and then casting to torch.int64 produces consistent results, which could help identify if the issue is related to integer rounding.
• Investigate if there are any open issues or discussions in the PyTorch repository related to this specific problem, as it may have been addressed or have a known workaround.

Example

import torch

# Explicitly round the results to the nearest integer
cpu = torch.round(torch.linspace(3.7, -3, 10, device="cpu")).int()
gpu = torch.round(torch.linspace(3.7, -3, 10, device="cuda")).int().cpu()

print("cpu:", cpu.tolist())
print("gpu:", gpu.tolist())

Notes

The provided code snippet and example are based on the information given in the issue and may not fully resolve the problem but should help in identifying a consistent rounding behavior.

Recommendation

Apply workaround: Use explicit rounding with torch.round to ensure consistent results between CPU and CUDA, as this approach directly addresses the suspected cause of the issue related to different rounding behaviors.