1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
|
#pragma once
#include <torch/csrc/distributed/rpc/message.h>
#include <torch/csrc/distributed/rpc/rpc_command_base.h>
#include <torch/csrc/distributed/rpc/types.h>
namespace torch {
namespace distributed {
namespace autograd {
// Internal system RPC to invoke distributed backward pass on remote nodes when
// 'rref.backward()' is invoked.
class TORCH_API RRefBackwardReq : public rpc::RpcCommandBase {
public:
RRefBackwardReq(
const rpc::RRefId& rrefId,
int64_t autogradContextId,
bool retainGraph = false);
const rpc::RRefId& getRRefId() const;
int64_t getAutogradContextId() const;
bool retainGraph() const;
// Serialization and deserialization methods.
c10::intrusive_ptr<rpc::Message> toMessageImpl() && override;
static std::unique_ptr<RRefBackwardReq> fromMessage(
const rpc::Message& message);
private:
const rpc::RRefId rrefId_;
const int64_t autogradContextId_;
const bool retainGraph_;
};
} // namespace autograd
} // namespace distributed
} // namespace torch
|
