提交 45e20602 编写于 作者: M Megvii Engine Team

feat(mgb/tensor): add proxy_to_comp_node for HostTensorND

GitOrigin-RevId: 8f03dbcbd110010008926019adca7a8854eb2607
上级 248d8bf0
......@@ -548,6 +548,15 @@ class TensorND {
ret.reset(storage().proxy_to_default_cpu(), layout());
return ret;
}
template<bool x = true, typename = std::enable_if_t<x && std::is_same<TensorStorage, HostTensorStorage>::value>>
HostTensorND proxy_to_comp_node(CompNode cn) const {
HostTensorStorage host_storage;
host_storage.reset(cn, m_storage.size(), m_storage.raw_storage());
HostTensorND ret;
ret.reset(host_storage, m_layout);
return ret;
}
};
/*!
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册