diff --git a/third_party/llga b/third_party/llga
index 5f36fafc3..1c0bff14e 160000
--- a/third_party/llga
+++ b/third_party/llga
@@ -1 +1 @@
-Subproject commit 5f36fafc3547b233f0f3c617ef917d729354aec3
+Subproject commit 1c0bff14e182864e835b1693db7d08c6f7e738b6
diff --git a/torch_ipex/csrc/jit/codegen/onednn/kernel.cpp b/torch_ipex/csrc/jit/codegen/onednn/kernel.cpp
index dd64d335b..91a437a74 100644
--- 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);