Add index to shape getter

wonjoo-wj · wonjoo-wj · commit 2ad82dea660b · 2022-05-17T22:46:08.000Z
diff --git a/torch_xla/csrc/ir.cpp b/torch_xla/csrc/ir.cpp
@@ -225,7 +225,7 @@ void ScopePusher::ResetScopes() { ResetScopeContext(); }
 
 const xla::Shape& GetXlaShape(const torch::lazy::Value& value) {
   XlaNode* casted = dynamic_cast<XlaNode*>(value.node.get());
-  return casted->xla_shape();
+  return casted->xla_shape(value.index);
 }
 
 }  // namespace torch_xla

Original file line number	Diff line number	Diff line change
`@@ -225,7 +225,7 @@ void ScopePusher::ResetScopes() { ResetScopeContext(); }`
`225`	`225`
`226`	`226`	`const xla::Shape& GetXlaShape(const torch::lazy::Value& value) {`
`227`	`227`	`XlaNode* casted = dynamic_cast<XlaNode*>(value.node.get());`
`228`		`- return casted->xla_shape();`
	`228`	`+ return casted->xla_shape(value.index);`
`229`	`229`	`}`
`230`	`230`
`231`	`231`	`} // namespace torch_xla`