mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-21 05:34:18 +08:00
Follows #136848 Pull Request resolved: https://github.com/pytorch/pytorch/pull/137072 Approved by: https://github.com/kwen2501
25 lines
736 B
C++
25 lines
736 B
C++
#include <ATen/record_function.h>
|
|
#include <torch/csrc/distributed/autograd/autograd.h>
|
|
|
|
namespace torch::distributed::autograd {
|
|
|
|
constexpr auto kDistAutogradBackwardProfilingKey =
|
|
"torch::distributed::autograd::backward";
|
|
|
|
void backward(
|
|
int64_t context_id,
|
|
const variable_list& roots,
|
|
bool retain_graph) {
|
|
C10_LOG_API_USAGE_ONCE("torch.distributed.autograd.backward");
|
|
RECORD_FUNCTION(
|
|
kDistAutogradBackwardProfilingKey, std::vector<c10::IValue>());
|
|
try {
|
|
DistEngine::getInstance().execute(context_id, roots, retain_graph);
|
|
} catch (std::exception& e) {
|
|
// FIXME: crashes if exception type is not RuntimeError
|
|
TORCH_CHECK(false, e.what());
|
|
}
|
|
}
|
|
|
|
} // namespace torch::distributed::autograd
|