索引CUDA张量列表报错，如何保留GPU存储完成批量索引？

问题核心本质说明 #

你遇到的两个报错是独立的：

• 第一个TypeError: only integer scalar arrays can be converted to a scalar index是Python原生列表的通用限制，和张量的存储设备无关：普通Python列表仅支持单个整数作为索引，不支持用索引列表/数组做批量取值。
• 第二个报错是CUDA设备上的张量无法直接转换为numpy数组，必须先调用.cpu()移动到主机内存才能转换，但该操作对你的索引需求完全不必要，不需要做CPU-GPU的来回数据拷贝。

可保留CUDA设备存储的解决方案 #

不需要将张量移动到CPU即可完成索引，推荐以下两种方案：

方案1：列表推导式（适配所有场景，零额外开销）

直接按索引列表遍历取值即可，所有张量全程保留在CUDA设备上，无任何数据拷贝：

X_train_fold = [X_train[i] for i in train_index]

该方法尤其适配样本长度不一致的场景（比如未做全局padding的NLP序列数据），拿到的子列表可直接用于后续GPU训练。

方案2：堆叠为统一CUDA张量（适配样本形状一致的场景，索引效率更高）

如果列表内所有张量的形状完全相同（比如已经做过全局padding的序列数据），可以直接用PyTorch的stack方法将列表转为单块CUDA张量，之后就可以像numpy数组一样直接用索引列表批量取值：

import torch
# 堆叠后的张量依然存储在cuda:0设备上
X_train_tensor = torch.stack(X_train)
# 直接批量索引，无设备转移
X_train_fold = X_train_tensor[train_index]

该方案后续的索引、切分、批处理效率更高，适合需要多次做数据划分的场景。

不推荐的操作 #

不要将CUDA张量转为numpy数组完成索引后再移回CUDA，该流程会产生大量不必要的跨设备数据拷贝，既浪费运算时间，也会增加内存和显存的额外开销。

该类场景的最优实践 #

• 若样本形状不一致，优先用列表推导式完成索引，全程无需修改张量的存储设备。
• 若样本形状一致，优先将列表堆叠为统一的CUDA张量，后续所有数据操作效率更高。
• 仅当需要使用numpy专属的运算能力时，才将对应张量移动到CPU转numpy，普通的数据切分、索引操作不需要跨设备转换。

内容的提问来源于stack exchange，提问作者Mahesha999