import torch
from torchvision.datasets import MNIST
from torchvision.transforms import ToTensor
# 加载 MNIST 数据集
mnist = MNIST(root='data', train=True, download=True, transform=ToTensor())
# 将灰度图像二值化
threshold = 128 # 设定阈值为 128
mnist_binary = torch.where(mnist.data >= threshold, torch.ones_like(mnist.data), torch.zeros_like(mnist.data))
# 输出二值化后的图像
print(mnist_binary[0])