原帖用裁剪后的 SqueezeNet 做预测，失败行是 output.data.numpy().argmax()：
TypeError: can’t convert CUDA tensor to numpy. Use Tensor.cpu() to copy the tensor to host memory first.
先决定结果要留在什么地方
原题的要求是“不把整个输出复制到 CPU”。可以在 GPU 上完成 argmax，得到类别索引；但若最终需要 Python 的整数或列表，仍会有少量数据传回主机，并可能触发同步。应减少不必要的全量复制，而不是承诺 Python 变量能够直接使用 CUDA 内存。
这行先把整个 CUDA 输出交给 NumPy，再计算最大值索引。NumPy 的普通 ndarray 使用 CPU 内存；若只想取预测类别，没有必要先复制全部输出。
在预测得到 output 后检查它的形状，确认类别轴：
print("josh050608")
print("output:", output.shape, output.device, output.dtype)
print("requires_grad:", output.requires_grad)
若输出是 [N,C]，按类别维取最大值：
predicted = output.argmax(dim=1)
# 单张图片，且 predicted 只有一个元素时：
index = predicted.item()
计算仍在 GPU 进行，只把最后的标量结果送回 Python。若是多张图片，应保留 predicted 张量，或在确实需要 CPU 结果时调用 predicted.cpu().tolist()。若输出是单个 [C] 向量，使用 output.argmax(dim=0).item()。
确实需要 NumPy 数组时，再用 array = output.detach().cpu().numpy()。detach() 处理梯度关系，.cpu() 处理设备，两者作用不同。不要继续通过 .data 绕开自动求导。
原帖没有给出输出形状，所以不能默认无参数的 argmax() 对多批图片也正确：它会把整个张量当成一个整体。验证时先用单张图片，再用两张图片确认返回两个类别，并检查索引范围在 [0,C)。
单张和一批图片，应该返回不同数量的索引
下面是独立的合成例子，有 CUDA 时在 GPU 上运行，否则在 CPU 上验证形状逻辑。两行分别代表两张图片，四列代表四类：
import torch
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
scores = torch.tensor([[0.1, 0.8, 0.2, -0.1],
[0.7, 0.2, 0.4, 0.9]], device=device)
indices = scores.argmax(dim=1)
assert indices.shape == (2,)
assert indices.device == scores.device
assert indices.cpu().tolist() == [1, 3]
print("batch indices:", indices)
single_index = scores[:1].argmax(dim=1).item()
assert single_index == 1
print("one image:", single_index)
如果直接写 scores.argmax()，得到的是把两张图片展平之后的一个位置，不能作为每张图片的预测类别。把这个例子替换回 SqueezeNet 时，先打印真实 output.shape，避免网络修改后还沿用旧轴编号。
| 后续用途 | 建议保留的结果 |
|---|---|
| 下一步继续做 GPU 张量运算 | output.argmax(dim=1)，仍为 GPU Tensor |
| 单张图片映射 Python 类名列表 | 单元素预测 .item() 后查类名 |
| 批量交给 CPU 评估工具 | 对预测索引 .cpu().tolist() |
| 必须导出全部分数 | output.detach().cpu().numpy() |
推理入口还应使用已经加载好的 model.eval()，并把前向放进 torch.no_grad()。这能减少推理中不必要的计算图，却不能让 NumPy 直接接管 CUDA 张量。验证时比较 PyTorch argmax 与 CPU 数组上按同一类别轴 argmax 的结果是否一致，再确认多图输入不会只返回一个全局位置。
参考：原始求助帖；技术依据：Tensor.numpy（2.5）、PyTorch 自动求导（2.5）、Tensor.item（2.5）。
