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
|
#include <torch/csrc/jit/passes/annotate_warns.h>
#include <atomic>
namespace torch {
namespace jit {
void AnnotateWarns(Block* b) {
static std::atomic<int64_t> idx(0);
for (Node* n : b->nodes()) {
for (Block* child_b : n->blocks()) {
AnnotateWarns(child_b);
}
if (n->kind() != aten::warn) {
continue;
}
n->i_(attr::warn_id, idx);
idx++;
}
}
void AnnotateWarns(const std::shared_ptr<Graph>& graph) {
AnnotateWarns(graph->block());
}
} // namespace jit
} // namespace torch
|
