LibTorch中TensorOptions的使用

*Major*

已于 2022-08-27 19:00:01 修改

阅读量1.2k

点赞数

文章标签：算法

于 2022-08-27 18:58:42 首次发布

本文链接：https://blog.csdn.net/qq_41375318/article/details/126561191

版权

$L ib T orc h 中 T e n sor Opt i o n s 的使用$

TensorOptions是用来配置Tensor属性用的

auto options =  torch::TensorOptions()
    .dtype(torch::kFloat32)
    .layout(torch::kStrided)
    .device(torch::kCUDA, 1)
    .requires_grad(true);

torch::Tensor tensor = torch::full({3, 4}, /*value=*/123, options);

assert(tensor.dtype() == torch::kFloat32);
assert(tensor.layout() == torch::kStrided);
assert(tensor.device().type() == torch::kCUDA); // or device().is_cuda()
assert(tensor.device().index() == 1);
assert(tensor.requires_grad());

示例：

#include <torch/torch.h>
using namespace std;
int main()
{

    auto options = torch::TensorOptions()
        .dtype(torch::kFloat32)
        .requires_grad(true);

    torch::Tensor w = torch::randn({ 1 }, options);

    assert(w.dtype() == torch::kFloat32);
    assert(w.requires_grad());

    cout << (w.dtype() == torch::kFloat32) << "  " << w.requires_grad() << endl;

	system("pause");
	return 0;
}