帮我写一段生成训练集与测试集的代码

当生成训练集和测试集时，通常需要将数据集分割成两个部分。以下是一个示例代码，用于生成训练集和测试集：

import torch
from sklearn.model_selection import train_test_split

# 假设有特征数据 features 和标签数据 labels

# 使用 train_test_split 函数将数据集分割成训练集和测试集
train_features, test_features, train_labels, test_labels = train_test_split(features, labels, test_size=0.2)

# 将数据转换为 PyTorch 张量
train_features = torch.tensor(train_features)
train_labels = torch.tensor(train_labels)
test_features = torch.tensor(test_features)
test_labels = torch.tensor(test_labels)

在这个示例中，我们使用了`train_test_split`函数从`features`和`labels`中分割出训练集和测试集。`test_size=0.2`表示将20%的数据作为测试集，80%的数据作为训练集。分割后的训练集特征数据存储在`train_features`中，训练集标签数据存储在`train_labels`中，测试集特征数据存储在`test_features`中，测试集标签数据存储在`test_labels`中。

最后，我们将这些数据转换为PyTorch张量，以便在PyTorch模型中使用。请注意，在实际应用中，你可能需要根据你的数据集格式进行适当的修改。