Skip to content

Commit 74df712

Browse files
authored
fix: handle device in the same way as dtype in aten.full_like decomposition (#3538)
1 parent 8ccf325 commit 74df712

File tree

1 file changed

+1
-1
lines changed

1 file changed

+1
-1
lines changed

py/torch_tensorrt/dynamo/lowering/_decompositions.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -434,7 +434,7 @@ def full_like_decomposition(*args, **kwargs) -> torch.Tensor:
434434
shape = args[0].shape
435435
fill_value = args[1]
436436
kwargs["dtype"] = kwargs.get("dtype", None) or input.dtype
437-
kwargs["device"] = input.device
437+
kwargs["device"] = kwargs.get("device", None) or input.device
438438
return torch.full(shape, fill_value, dtype=kwargs["dtype"], device=kwargs["device"])
439439

440440

0 commit comments

Comments
 (0)