Skip to content

Commit

Permalink
refactor(//core/lowering): Update lowering for PyTorch 1.6.0
Browse files Browse the repository at this point in the history
Signed-off-by: Naren Dasan <naren@narendasan.com>
Signed-off-by: Naren Dasan <narens@nvidia.com>
  • Loading branch information
narendasan committed Jul 24, 2020
1 parent fe06d09 commit 378690b
Showing 1 changed file with 7 additions and 4 deletions.
11 changes: 7 additions & 4 deletions core/lowering/lowering.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,6 @@
#include "torch/csrc/jit/passes/lower_graph.h"
#include "torch/csrc/jit/passes/lower_tuples.h"
#include "torch/csrc/jit/passes/peephole.h"
#include "torch/csrc/jit/passes/quantization.h"

#include "core/util/prelude.h"
#include "core/lowering/lowering.h"
Expand Down Expand Up @@ -50,8 +49,7 @@ torch::jit::Module LowerModule(const torch::jit::script::Module& mod) {
return mod_;
}

std::pair<std::shared_ptr<torch::jit::Graph>, std::vector<at::Tensor>> Lower(const torch::jit::script::Module& mod,
std::string method_name) {
std::pair<std::shared_ptr<torch::jit::Graph>, std::vector<at::Tensor>> Lower(const torch::jit::script::Module& mod, std::string method_name) {
auto lowered_mod = LowerModule(mod);
auto g = lowered_mod.get_method(method_name).graph();
LOG_GRAPH(*g);
Expand All @@ -62,9 +60,14 @@ std::pair<std::shared_ptr<torch::jit::Graph>, std::vector<at::Tensor>> Lower(con
lowering::LowerGraph(g);
//=[torch::jit::FoldConvBatchNorm2d(lowered_mod);
LOG_GRAPH("LibTorch Lowering");
auto graph_and_parameters = torch::jit::LowerGraph(*g, lowered_mod._ivalue());
auto graph_and_ivalues = torch::jit::LowerGraph(*g, lowered_mod._ivalue());
// Is this necessary?
lowering::LowerBlock(g->block());
std::pair<std::shared_ptr<torch::jit::Graph>, std::vector<at::Tensor>> graph_and_parameters;
for (auto i : graph_and_ivalues.second) {
graph_and_parameters.second.push_back(i.toTensor());
}
graph_and_parameters.first = graph_and_ivalues.first;
return graph_and_parameters;
}

Expand Down

0 comments on commit 378690b

Please sign in to comment.