Convert tensor types in `make_tensor_view`. (#2755)

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

1 files changed, 5 insertions, 2 deletions

diff --git a/prelude/slang-torch-prelude.h b/prelude/slang-torch-prelude.h
index 4844e9248..70c516a3a 100644
--- a/prelude/slang-torch-prelude.h
+++ b/prelude/slang-torch-prelude.h
@@ -68,11 +68,14 @@ struct CudaTaskMemoryAllocator
     }
 };
 
-TensorView make_tensor_view(CudaTaskMemoryAllocator* allocator, torch::Tensor val, const char* name)
+TensorView make_tensor_view(CudaTaskMemoryAllocator* allocator, torch::Tensor val, const char* name, torch::ScalarType targetScalarType)
 {
+    // Convert device and scalar types.
     if (!val.device().is_cuda())
         val = val.to(torch::kCUDA);
- 
+    if (val.dtype() != targetScalarType)
+        val = val.to(targetScalarType);
+
     TensorView res = {};
     res.dimensionCount = val.dim();
     res.strides = allocator->allocUIntArray(val.dim());