TL;DR

The issue can be fixed by changing the data type of the tensor from torch.uint8 to a larger integer type, such as torch.int32 or torch.int64, to prevent overflow when calculating the sum.

Guidance

• The problem is likely caused by integer overflow when calculating the sum of the tensor elements, as the sum of the elements in the example exceeds the maximum value that can be represented by torch.uint8 (255).
• To verify this, try printing the sum of the elements in the tensor without specifying the data type, and compare the result with the expected sum.
• To mitigate this issue, change the data type of the tensor to a larger integer type, such as torch.int32 or torch.int64, when creating the tensor.
• Additionally, consider using the torch.sum function without specifying the data type, and let PyTorch automatically determine the data type of the result.

Example

x_mps = torch.tensor(example, dtype=torch.int32, device="mps")
x_cpu = torch.tensor(example, dtype=torch.int32, device="cpu")

y_cpu = x_cpu.sum().item()
y_mps = x_mps.sum().item()

print(f"{y_cpu=}")
print(f"{y_mps=}")

Notes

The issue is specific to the torch.uint8 data type and the MPS device, and may not occur with other data types or devices.

Recommendation

Apply workaround by changing the data type of the tensor to a larger integer type, such as torch.int32 or torch.int64, to prevent overflow when calculating the sum. This is because the sum of the elements in the example exceeds the maximum value that can be represented by torch.uint8 (255).