【pytorch】torch.utils.data.TensorDataset()原版与新版的差异

使用时发现出错了，如下：

原因是新版把之前的data_tensor 和target_tensor去掉了，输入变成了可变参数，也就是我们平常使用*args

class TensorDataset(Dataset):
    """Dataset wrapping tensors.

    Each sample will be retrieved by indexing tensors along the first dimension.

    Arguments:
        *tensors (Tensor): tensors that have the same size of the first dimension.
    """

    def __init__(self, *tensors):
        assert all(tensors[0].size(0) == tensor.size(0) for tensor in tensors)
        self.tensors = tensors

    def __getitem__(self, index):
        return tuple(tensor[index] for tensor in self.tensors)

    def __len__(self):
        return self.tensors[0].size(0)

所以新版的使用方法是直接传入参数：

# 原版使用方法
train_dataset = Data.TensorDataset(data_tensor=x, target_tensor=y)

# 新版使用方法
train_dataset = Data.TensorDataset(x,y)