pytorch Tensor.to(device)和model.to(device) 的区别

康康同学97

于 2021-06-10 16:40:21 发布

阅读量6.9k

点赞数 13

分类专栏： pytorh学习记录文章标签： pytorch cuda 深度学习

本文链接：https://blog.csdn.net/qq_38436266/article/details/117784900

版权

3 篇文章

订阅专栏

本文详细介绍了在PyTorch框架中如何正确地将Module对象及Tensor数据转移到GPU进行加速计算。对于Module对象，只需直接调用to(device)即可；而Tensor数据则需接收返回值来确保正确设置device。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

使用GPU训练的时候，需要将Module对象和Tensor类型的数据送入到device。通常会使用 to.(device)。但是需要注意的是：

来自pytorch官方文档的说明：

Tensor.to(device)

Module.to(device)

# Module对象设置device的写法
model.to(device)

# Tensor类型的数据设置 device 的写法。
samples = samples.to(device)