如何在torch中获取箭图的数据？

import torch
from torchvision import datasets, transforms

# 加载箭图数据集
# 这里以MNIST数据集为例
train_dataset = datasets.MNIST(root='./data', train=True, download=True, transform=transforms.ToTensor())
test_dataset = datasets.MNIST(root='./data', train=False, download=True, transform=transforms.ToTensor())

# 创建数据加载器
batch_size = 64
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=batch_size, shuffle=False)

# 迭代获取箭图数据
for images, labels in train_loader:
    # 在这里可以对箭图数据进行处理
    # images是一个大小为[batch_size, 1, 28, 28]的张量，表示一批箭图数据
    # labels是一个大小为[batch_size]的张量，表示一批箭图数据对应的标签
    # 可以根据具体需求进行操作
    pass