Skip to content

Commit

Permalink
fix(//core/conversion): Add special case for If and Loop
Browse files Browse the repository at this point in the history
which are handled by conversion and not a converter

Signed-off-by: Naren Dasan <naren@narendasan.com>
Signed-off-by: Naren Dasan <narens@nvidia.com>
  • Loading branch information
narendasan committed Apr 14, 2022
1 parent 3c59ece commit eacde8d
Showing 1 changed file with 4 additions and 1 deletion.
5 changes: 4 additions & 1 deletion core/conversion/conversion.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,10 @@ namespace conversion {
bool isNodeConversionIgnored(const torch::jit::Node* n);

bool OpSupported(const torch::jit::Node* n) {
return evaluators::shouldEvalAtConversionTime(n) || converters::node_is_convertable(n);
return evaluators::shouldEvalAtConversionTime(n)
|| converters::node_is_convertable(n)
|| n->kind() == torch::jit::prim::Loop
|| n->kind() == torch::jit::prim::If;
}

c10::optional<torch::jit::IValue> EvaluateNode(ConversionCtx* ctx, const torch::jit::Node* n, int level, int limit) {
Expand Down

0 comments on commit eacde8d

Please sign in to comment.