revert onednn update (intel#189)

* Revert "[LLGA] remove infer_shape due to backend change (intel#176)" This reverts commit c1fe854. * Revert "[LLGA] ww35 update backend (intel#175)" This reverts commit eccafbe.
tenzen-y · Sep 2, 2021 · f4af2b9 · f4af2b9
1 parent c1fe854
commit f4af2b9
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 2 deletions.
diff --git a/third_party/llga b/third_party/llga
diff --git a/torch_ipex/csrc/jit/codegen/onednn/kernel.cpp b/torch_ipex/csrc/jit/codegen/onednn/kernel.cpp
@@ -70,10 +70,14 @@ ArgSpecs LlgaKernel::specializeInputSpecs(const TensorArgs &inputs) const {
 
 ArgSpecs LlgaKernel::specializeOutputSpecs(const partition &partition,
                                            const ArgSpecs &inputSpecs) const {
+  auto inputs = fmap(inputSpecs, toLogicalTensor);
+  auto outputs = fmap(graph_->outputs(), toLogicalTensor);
+  partition.infer_shape(inputs, outputs);
+
   ArgSpecs outputSpecs;
   outputSpecs.reserve(nOutputs_);
   for (size_t i = 0; i < nOutputs_; i++) {
-    auto spec = ArgSpec(graph_->outputs()[i]);
+    auto spec = ArgSpec(outputs[i]);
 
     if (spec.is_quantized())
       spec = getQuantizedSpec(spec, i);