Small fixes to TorchTensor. (#2790)

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

3 files changed, 12 insertions, 3 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("))");
diff --git a/source/slang/slang-ir-insts.h b/source/slang/slang-ir-insts.h
index 26d00d4df..a19896287 100644
--- a/source/slang/slang-ir-insts.h
+++ b/source/slang/slang-ir-insts.h
@@ -2807,7 +2807,7 @@ public:
 
     IRArrayListType* getArrayListType(IRType* elementType);
     IRTensorViewType* getTensorViewType(IRType* elementType);
-    IRTorchTensorType* getTorchTensorType();
+    IRTorchTensorType* getTorchTensorType(IRType* elementType);
 
     IRDifferentialPairType* getDifferentialPairType(
         IRType* valueType,
diff --git a/source/slang/slang-ir.cpp b/source/slang/slang-ir.cpp
index 6ca05d2d6..bd2271953 100644
--- a/source/slang/slang-ir.cpp
+++ b/source/slang/slang-ir.cpp
@@ -2837,9 +2837,9 @@ namespace Slang
             (IRInst**)&elementType);
     }
 
-    IRTorchTensorType* IRBuilder::getTorchTensorType()
+    IRTorchTensorType* IRBuilder::getTorchTensorType(IRType* elementType)
     {
-        return (IRTorchTensorType*)getType(kIROp_TorchTensorType, 0, nullptr);
+        return (IRTorchTensorType*)getType(kIROp_TorchTensorType, 1, (IRInst**)&elementType);
     }
 
     IRDifferentialPairType* IRBuilder::getDifferentialPairType(