fix: cumsum add_constant bug fix (add dtype for np zeros)

pytorch · Oct 22, 2024 · 123e2a8 · 123e2a8
1 parent 40193dc
commit 123e2a8
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/py/torch_tensorrt/dynamo/conversion/impl/slice/ops.py b/py/torch_tensorrt/dynamo/conversion/impl/slice/ops.py
@@ -370,7 +370,7 @@ def cumsum(
         )
     else:
         new_dims = tuple(data.shape)
-        zeros = np.zeros(new_dims)
+        zeros = np.zeros(new_dims, dtype=np.float32)
         zero_trttensor = get_trt_tensor(ctx, zeros, f"{name}_initial_value")
 
     running_sum = loop.add_recurrence(zero_trttensor)