Pytork的二进制输入

我可以问一下,如何使pytorch中加载的数据在加载后变成二进制? 就像Tensorflow可以通过以下方式实现:

train_data  = mnist.input_data.read_data_sets(data_directory, one_hot=True)

pytorch如何实现 one_hot=True 效应

我现在拥有的data_loader是:

torch.set_default_tensor_type('torch.FloatTensor')
train_loader = torch.utils.data.DataLoader(
        datasets.MNIST('data/', train=True, download=True,
                       transform=transforms.Compose([
                         # transforms.RandomHorizontalFlip(),
                           transforms.ToTensor()])),
        batch_size=batch_size, shuffle=False)

我想将train_loader中的数据进行二值化。 现在我要做的是:加载数据后,

for data,_ in train_loader:
    torch.round(data) 
    data = Variable(data)

使用 torch.round() 作用这是正确的吗?