1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
|
#include <torch/csrc/jit/passes/remove_expands.h>
namespace torch::jit {
static void RemoveExpands(Block* block) {
for (auto it = block->nodes().begin(), end = block->nodes().end(); it != end;
++it) {
for (auto sub : it->blocks())
RemoveExpands(sub);
if (it->kind() == aten::expand && it->get<bool>(attr::implicit) == true) {
it->output()->replaceAllUsesWith(it->namedInput(attr::self));
it.destroyCurrent();
}
}
}
void RemoveExpands(const std::shared_ptr<Graph>& graph) {
RemoveExpands(graph->block());
}
} // namespace torch::jit
|
