Small fixes to TorchTensor. (#2790)

Co-authored-by: Yong He <yhe@nvidia.com>

1 files changed, 9 insertions, 0 deletions

diff --git a/source/slang/slang-emit-torch.cpp b/source/slang/slang-emit-torch.cpp
index c198b011d..819a6a136 100644
--- a/source/slang/slang-emit-torch.cpp
+++ b/source/slang/slang-emit-torch.cpp
@@ -165,6 +165,15 @@ bool TorchCppSourceEmitter::tryEmitInstExprImpl(IRInst* inst, const EmitOpInfo&
                 auto arg = inst->getOperand(i);
                 emitOperand(arg, getInfo(EmitOp::General));
             }
+            if (as<IRTorchTensorType>(inst->getDataType()))
+            {
+                if (auto vectorType = as<IRVectorType>(inst->getDataType()->getOperand(0)))
+                {
+                    // If the element type of the tensor is a vector, we need to add the vector size to the shape.
+                    m_writer->emit(", ");
+                    emitOperand(vectorType->getElementCount(), getInfo(EmitOp::General));
+                }
+            }
             m_writer->emit("}, torch::TensorOptions().device(torch::kCUDA).dtype(");
             emitTorchScalarTypeName(m_writer, inst->getDataType());
             m_writer->emit("))");