mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-20 21:14:14 +08:00
ab2ce3c50e
Related ISSUE: https://github.com/pytorch/pytorch/issues/148114 Pull Request resolved: https://github.com/pytorch/pytorch/pull/163264 Approved by: https://github.com/albanD, https://github.com/cyyever
50 lines
1.3 KiB
C++
50 lines
1.3 KiB
C++
#pragma once
|
|
|
|
#include <ATen/core/Tensor.h>
|
|
#include <torch/csrc/python_headers.h>
|
|
|
|
#include <torch/csrc/utils/python_arg_parser.h>
|
|
|
|
namespace torch::autograd::utils {
|
|
|
|
// The parameter allow_copy is to accept copy for Tensor.to (and by proxy
|
|
// PackedSequences.to) but not nn.Module.to.
|
|
inline std::tuple<
|
|
std::optional<at::Device>,
|
|
std::optional<at::ScalarType>,
|
|
bool,
|
|
bool,
|
|
std::optional<at::MemoryFormat>>
|
|
parse_to_conversion(PythonArgs& r, bool allow_copy) {
|
|
if (r.idx == 0) {
|
|
TORCH_CHECK(
|
|
allow_copy || r.isNone(3), ".to() does not accept copy argument");
|
|
return std::make_tuple(
|
|
r.deviceOptional(0),
|
|
r.scalartypeOptional(1),
|
|
r.toBool(2),
|
|
r.toBool(3),
|
|
r.memoryformatOptional(4));
|
|
} else if (r.idx == 1) {
|
|
TORCH_CHECK(
|
|
allow_copy || r.isNone(2), ".to() does not accept copy argument");
|
|
return std::make_tuple(
|
|
std::nullopt,
|
|
r.scalartype(0),
|
|
r.toBool(1),
|
|
r.toBool(2),
|
|
r.memoryformatOptional(3));
|
|
} else {
|
|
auto tensor = r.tensor(0);
|
|
TORCH_CHECK(
|
|
allow_copy || r.isNone(2), ".to() does not accept copy argument");
|
|
return std::make_tuple(
|
|
tensor.device(),
|
|
tensor.scalar_type(),
|
|
r.toBool(1),
|
|
r.toBool(2),
|
|
r.memoryformatOptional(3));
|
|
}
|
|
}
|
|
} // namespace torch::autograd::utils
|